Skip to main content

bamboo_engine/
llm_summarizer.rs

1//! LLM-backed conversation summarizer.
2//!
3//! `LlmSummarizer` is the infrastructure-coupled implementation of
4//! `bamboo_compression::Summarizer`: it calls the session model to produce a
5//! rich summary of compressed/removed messages. Callers without an explicit
6//! model budget use a conservative compatibility budget, so every LLM-backed
7//! path still runs through bounded map/reduce. Compatibility callers may fall
8//! back to the pure `HeuristicSummarizer`; explicitly budgeted production passes
9//! surface every failed stage so callers can preserve session state atomically.
10//! It lives in the engine (not in bamboo-compression) so that the compression
11//! crate stays free of any LLM-provider dependency.
12
13use std::sync::Arc;
14
15use async_trait::async_trait;
16use futures::StreamExt;
17
18use bamboo_compression::{
19    HeuristicSummarizer, MessageSegmenter, Summarizer, TiktokenTokenCounter, TokenBudget,
20    TokenCounter,
21};
22use bamboo_domain::ReasoningEffort;
23use bamboo_domain::{
24    ContextBlock, ContextBlockPriority, ContextBlockStability, ContextBlockType, Message, Role,
25};
26use bamboo_llm::LLMChunk;
27use bamboo_llm::{LLMProvider, LLMRequestOptions};
28
29const COMPATIBILITY_CONTEXT_WINDOW_TOKENS: u32 = 32_000;
30const COMPATIBILITY_MAX_OUTPUT_TOKENS: u32 = 8_000;
31const COMPATIBILITY_SAFETY_MARGIN_TOKENS: u32 = 1_000;
32const COMPATIBILITY_SAFE_WINDOW_PERCENT: u8 = 80;
33
34/// Mode controlling how the LLM summarizer handles existing summaries.
35#[derive(Debug, Clone, Default)]
36pub enum SummaryMode {
37    /// Generate a complete summary from scratch (default).
38    #[default]
39    FullRewrite,
40    /// Update an existing summary by incorporating new information incrementally.
41    IncrementalMerge,
42}
43
44/// Hard request limits for a logical compression pass.
45#[derive(Debug, Clone)]
46pub struct SummaryRequestBudget {
47    pub context_window_tokens: u32,
48    pub max_output_tokens: u32,
49    pub safety_margin_tokens: u32,
50    pub safe_window_percent: u8,
51    pub target_summary_tokens: u32,
52    pub target_ratio: f64,
53}
54
55impl SummaryRequestBudget {
56    pub fn from_token_budget(
57        budget: &TokenBudget,
58        safe_window_percent: u8,
59        target_summary_tokens: u32,
60        target_ratio: f64,
61    ) -> Self {
62        Self {
63            context_window_tokens: budget.max_context_tokens.max(1),
64            max_output_tokens: budget.max_output_tokens.max(1),
65            safety_margin_tokens: budget.safety_margin,
66            safe_window_percent: safe_window_percent.clamp(10, 95),
67            target_summary_tokens: target_summary_tokens.max(1),
68            target_ratio: if target_ratio.is_finite() && target_ratio > 0.0 {
69                target_ratio.clamp(0.01, 0.50)
70            } else {
71                0.20
72            },
73        }
74    }
75
76    fn safe_request_tokens(&self) -> u32 {
77        self.context_window_tokens
78            .saturating_mul(self.safe_window_percent as u32)
79            .saturating_div(100)
80            .max(1)
81    }
82}
83
84#[derive(Debug, Clone)]
85pub struct SummarizationReport {
86    pub content: String,
87    pub represented_source_tokens: u32,
88    pub target_summary_tokens: u32,
89    pub actual_summary_tokens: u32,
90    pub map_calls: u32,
91    pub reduce_calls: u32,
92    pub fallback_used: bool,
93    pub budget_clamped: bool,
94    pub budget_clamp_reason: Option<String>,
95}
96
97#[derive(Debug, Clone)]
98pub struct SummarizationProgress {
99    pub stage: String,
100    pub stage_index: usize,
101    pub stage_count: usize,
102    pub estimated_input_tokens: u32,
103    pub requested_output_tokens: u32,
104    pub safe_request_tokens: u32,
105    pub model_context_tokens: u32,
106}
107
108type SummarizationProgressCallback = dyn Fn(&SummarizationProgress) + Send + Sync;
109
110#[derive(Debug, Clone)]
111struct SourceUnit {
112    text: String,
113    represented_source_tokens: u32,
114    first_message_id: String,
115    last_message_id: String,
116    continuation_part: usize,
117}
118
119#[derive(Debug, Clone)]
120struct SummaryPart {
121    content: String,
122    represented_source_tokens: u32,
123    first_message_id: String,
124    last_message_id: String,
125}
126
127/// LLM-based summarizer that calls the current session's model to generate
128/// a rich summary of compressed/removed messages.
129///
130/// Calls without an explicit request budget fall back to
131/// [`HeuristicSummarizer`] if the bounded LLM pipeline fails, UNLESS
132/// [`with_heuristic_fallback_on_error(false)`](Self::with_heuristic_fallback_on_error)
133/// is set. Budgeted calls always surface failed map/reduce stages so the archive
134/// transaction cannot commit a heuristic summary that represents a failed pass.
135pub struct LlmSummarizer {
136    llm: Arc<dyn LLMProvider>,
137    model: String,
138    /// Optional existing summary to build upon (incremental summarization).
139    existing_summary: Option<String>,
140    /// Structured runtime context blocks that should inform summarization.
141    context_blocks: Vec<ContextBlock>,
142    /// Optional user-provided instructions that override/extend the default summary focus.
143    custom_instructions: Option<String>,
144    /// Controls how the summarizer handles existing summaries.
145    summary_mode: SummaryMode,
146    /// When true (default), a transient compatibility-path LLM failure is
147    /// recovered by falling back to the [`HeuristicSummarizer`]. Explicitly
148    /// budgeted calls always return stage failures so a multi-request pass
149    /// remains atomic. The empty-response fallback is unaffected either way.
150    /// (issues #238, #763)
151    heuristic_fallback_on_error: bool,
152    /// Exact selected summarization-model limits and source-derived output goal.
153    request_budget: Option<SummaryRequestBudget>,
154    /// Correlates every map/reduce request with the eventual persisted
155    /// compression event (or with a failed logical pass).
156    logical_pass_id: Option<String>,
157    logical_phase: Option<String>,
158    progress_callback: Option<Arc<SummarizationProgressCallback>>,
159}
160
161impl LlmSummarizer {
162    pub fn new(
163        llm: Arc<dyn LLMProvider>,
164        model: String,
165        existing_summary: Option<String>,
166        task_list_prompt: Option<String>,
167    ) -> Self {
168        let context_blocks = task_list_prompt
169            .as_deref()
170            .map(str::trim)
171            .filter(|value| !value.is_empty())
172            .map(|task_list| {
173                vec![ContextBlock::new(
174                    ContextBlockType::TaskSnapshot,
175                    ContextBlockPriority::High,
176                    ContextBlockStability::RoundDynamic,
177                    "Current Task List",
178                    task_list,
179                )]
180            })
181            .unwrap_or_default();
182
183        Self {
184            llm,
185            model,
186            existing_summary,
187            context_blocks,
188            custom_instructions: None,
189            summary_mode: SummaryMode::default(),
190            heuristic_fallback_on_error: true,
191            request_budget: None,
192            logical_pass_id: None,
193            logical_phase: None,
194            progress_callback: None,
195        }
196    }
197
198    /// Control whether a transient compatibility-path LLM *error* recovers via
199    /// the heuristic summarizer (default `true`) or surfaces to the caller
200    /// (`false`). Explicitly budgeted map/reduce failures always surface.
201    /// (issues #238, #763)
202    pub fn with_heuristic_fallback_on_error(mut self, enabled: bool) -> Self {
203        self.heuristic_fallback_on_error = enabled;
204        self
205    }
206
207    pub fn with_context_blocks(mut self, context_blocks: Vec<ContextBlock>) -> Self {
208        self.context_blocks = context_blocks;
209        self
210    }
211
212    pub fn with_custom_instructions(mut self, instructions: Option<String>) -> Self {
213        self.custom_instructions = instructions;
214        self
215    }
216
217    pub fn with_summary_mode(mut self, mode: SummaryMode) -> Self {
218        self.summary_mode = mode;
219        self
220    }
221
222    pub fn with_request_budget(mut self, budget: SummaryRequestBudget) -> Self {
223        self.request_budget = Some(budget);
224        self
225    }
226
227    pub fn with_logical_pass_context(
228        mut self,
229        logical_pass_id: impl Into<String>,
230        logical_phase: impl Into<String>,
231    ) -> Self {
232        self.logical_pass_id = Some(logical_pass_id.into());
233        self.logical_phase = Some(logical_phase.into());
234        self
235    }
236
237    pub fn with_progress_callback(mut self, callback: Arc<SummarizationProgressCallback>) -> Self {
238        self.progress_callback = Some(callback);
239        self
240    }
241
242    fn append_shared_context(&self, user_content: &mut String) {
243        if let Some(ref existing) = self.existing_summary {
244            user_content.push_str("## Previous Summary\n\n");
245            user_content.push_str(existing);
246            user_content.push_str("\n\n---\n\n");
247        }
248
249        if !self.context_blocks.is_empty() {
250            user_content.push_str("## Compression Context Blocks\n\n");
251            for block in &self.context_blocks {
252                user_content.push_str(&format!(
253                    "### {}\n- type: {}\n- priority: {}\n- stability: {}\n\n{}\n\n",
254                    block.title.trim(),
255                    block.block_type.as_str(),
256                    block.priority.as_str(),
257                    block.stability.as_str(),
258                    block.content.trim(),
259                ));
260            }
261            user_content.push_str("---\n\n");
262        }
263
264        if let Some(ref instructions) = self.custom_instructions {
265            if !instructions.trim().is_empty() {
266                user_content.push_str("## Custom Compression Instructions\n\n");
267                user_content.push_str(instructions.trim());
268                user_content.push_str("\n\n---\n\n");
269            }
270        }
271    }
272
273    fn render_shared_context(&self) -> String {
274        let mut content = String::new();
275        self.append_shared_context(&mut content);
276        content
277    }
278
279    fn render_message_block(message: &Message) -> Option<String> {
280        let role_label = match message.role {
281            Role::User => "User",
282            Role::Assistant => "Assistant",
283            Role::Tool => "Tool Result",
284            Role::System => return None,
285        };
286        let mut block = String::new();
287        if let Some(ref tool_calls) = message.tool_calls {
288            if !tool_calls.is_empty() {
289                let tool_names = tool_calls
290                    .iter()
291                    .map(|call| call.function.name.as_str())
292                    .collect::<Vec<_>>();
293                block.push_str(&format!(
294                    "**{}** [message_id: {}; called tools: {}]:\n",
295                    role_label,
296                    message.id,
297                    tool_names.join(", ")
298                ));
299            } else {
300                block.push_str(&format!(
301                    "**{}** [message_id: {}]:\n",
302                    role_label, message.id
303                ));
304            }
305        } else {
306            block.push_str(&format!(
307                "**{}** [message_id: {}]:\n",
308                role_label, message.id
309            ));
310        }
311        if let Some(ref tool_call_id) = message.tool_call_id {
312            block.push_str(&format!("(tool_call_id: {})\n", tool_call_id));
313        }
314        block.push_str(&message.content);
315        block.push_str("\n\n");
316        Some(block)
317    }
318
319    fn render_messages(messages: &[Message]) -> String {
320        messages
321            .iter()
322            .filter_map(Self::render_message_block)
323            .collect::<String>()
324    }
325
326    fn build_map_messages(&self, units: &[SourceUnit], target_tokens: u32) -> Vec<Message> {
327        let represented_source_tokens = units.iter().fold(0u32, |total, unit| {
328            total.saturating_add(unit.represented_source_tokens)
329        });
330        let mut user_content = format!(
331            "Summarize this chronological source slice as loss-aware working memory.\n\
332             It represents approximately {represented_source_tokens} raw source tokens.\n\
333             Target approximately {target_tokens} output tokens. Preserve concrete requirements, \
334             decisions, paths, commands, errors, test results, tool outcomes, active work, and next \
335             steps. Do not turn it into a tiny high-level synopsis.\n\n## Source Slice\n\n"
336        );
337        for unit in units {
338            user_content.push_str(&format!(
339                "### Source range {}..{} (continuation part {})\n\n",
340                unit.first_message_id, unit.last_message_id, unit.continuation_part
341            ));
342            user_content.push_str(&unit.text);
343            user_content.push('\n');
344        }
345        user_content.push_str(
346            "\nReturn only the chronological partial summary. Do not claim this is the final conversation summary.",
347        );
348        vec![
349            Message::system(
350                "You are the map stage of a bounded conversation-compression pipeline. \
351                 Preserve detailed facts and ordering for a later reducer.",
352            ),
353            Message::user(user_content),
354        ]
355    }
356
357    fn build_reduce_messages(
358        &self,
359        parts: &[SummaryPart],
360        target_tokens: u32,
361        include_shared_context: bool,
362    ) -> Vec<Message> {
363        let represented_source_tokens = parts.iter().fold(0u32, |total, part| {
364            total.saturating_add(part.represented_source_tokens)
365        });
366        let mut user_content = String::new();
367        if include_shared_context {
368            self.append_shared_context(&mut user_content);
369        }
370        user_content.push_str(&format!(
371            "## Summary Size Budget\nTarget approximately {target_tokens} tokens. The partials \
372             below represent approximately {represented_source_tokens} raw source tokens. Derive \
373             detail from represented source size, not from the already-compressed partial length; \
374             do not apply the target ratio again.\n\n## Ordered Partial Summaries\n\n"
375        ));
376        for (index, part) in parts.iter().enumerate() {
377            user_content.push_str(&format!(
378                "### Partial {} — source {}..{} ({} represented raw tokens)\n\n{}\n\n",
379                index + 1,
380                part.first_message_id,
381                part.last_message_id,
382                part.represented_source_tokens,
383                part.content.trim(),
384            ));
385        }
386        if include_shared_context {
387            user_content.push_str(
388                "## Required Final Sections\n1. Pre-compression in-flight work\n2. Current active objective\n3. Requirement checklist with status and 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",
389            );
390        }
391        user_content.push_str(
392            "Return only the merged summary. Preserve source chronology and remove only genuine duplication.",
393        );
394        let system_prompt = if include_shared_context {
395            match self.summary_mode {
396                SummaryMode::FullRewrite => {
397                    "You are the final reduce stage of a bounded conversation-compression \
398                     pipeline. Merge all ordered partials and the supplied prior/runtime context \
399                     into one reliable working-memory summary."
400                }
401                SummaryMode::IncrementalMerge => {
402                    "You are the incremental final reduce stage of a bounded \
403                     conversation-compression pipeline. Update the supplied prior summary with the \
404                     ordered new partials and current runtime context. Newer facts supersede stale \
405                     prior state."
406                }
407            }
408        } else {
409            "You are an intermediate reduce stage of a bounded conversation-compression pipeline. \
410             Merge ordered partials without compounding the source-to-summary ratio."
411        };
412        vec![Message::system(system_prompt), Message::user(user_content)]
413    }
414
415    fn build_multipart_finalize_messages(
416        &self,
417        parts: &[SummaryPart],
418        target_tokens: u32,
419        shared_context_capsule: &str,
420        retain_shared_context_capsule: bool,
421    ) -> Vec<Message> {
422        let represented_source_tokens = parts.iter().fold(0u32, |total, part| {
423            total.saturating_add(part.represented_source_tokens)
424        });
425        let has_shared_context = !shared_context_capsule.trim().is_empty();
426        let mut user_content = String::new();
427        if has_shared_context {
428            let retention_instruction = if retain_shared_context_capsule {
429                "The shared-context capsule is retained once before the multipart updates, so do \
430                 not repeat unchanged capsule text; emit corrections, superseding facts, and \
431                 current runtime state when relevant."
432            } else {
433                "The shared-context capsule is reference material and is not stored separately. \
434                 Incorporate its durable prior state and current runtime facts wherever needed, \
435                 while applying its custom and hook-injected instructions as binding directives."
436            };
437            user_content.push_str(&format!(
438                "## Shared Context Capsule\n\n{shared_context_capsule}\n\n\
439                 Apply all custom and hook-injected instructions carried by the capsule. \
440                 {retention_instruction}\n\n"
441            ));
442        }
443        user_content.push_str(&format!(
444            "## Multipart Section Budget\n\n\
445             Target approximately {target_tokens} tokens for the source ranges in this section. \
446             They represent approximately {represented_source_tokens} raw source tokens. Preserve \
447             that source-derived allocation instead of applying the compression ratio again.\n\n\
448             ## Ordered Partial Summaries\n\n"
449        ));
450        for (index, part) in parts.iter().enumerate() {
451            user_content.push_str(&format!(
452                "### Partial {} — source {}..{} ({} represented raw tokens)\n\n{}\n\n",
453                index + 1,
454                part.first_message_id,
455                part.last_message_id,
456                part.represented_source_tokens,
457                part.content.trim(),
458            ));
459        }
460        user_content.push_str(
461            "Return only this finalized chronological multipart section. Preserve detailed evidence \
462             up to the allocated budget and make any newer correction explicit.",
463        );
464        vec![
465            Message::system(if has_shared_context {
466                "You are the instruction-aware multipart final stage of a bounded \
467                 conversation-compression pipeline. Finalize one ordered section under the supplied \
468                 shared context without compounding the source-to-summary ratio."
469            } else {
470                "You are the multipart final stage of a bounded conversation-compression pipeline. \
471                 Reduce one ordered section into its final form without compounding the \
472                 source-to-summary ratio."
473            }),
474            Message::user(user_content),
475        ]
476    }
477
478    fn target_for_source(&self, represented_source_tokens: u32) -> u32 {
479        let ratio = self
480            .request_budget
481            .as_ref()
482            .map(|budget| budget.target_ratio)
483            .unwrap_or(0.20);
484        ((represented_source_tokens as f64) * ratio).ceil().max(1.0) as u32
485    }
486
487    fn compatibility_request_budget(&self, messages: &[Message]) -> SummaryRequestBudget {
488        let counter = TiktokenTokenCounter::default();
489        let represented_source_tokens = counter.count_messages(messages);
490        let previous_summary_tokens = self
491            .existing_summary
492            .as_deref()
493            .map(|summary| counter.count_text(summary))
494            .unwrap_or(0);
495        SummaryRequestBudget {
496            context_window_tokens: COMPATIBILITY_CONTEXT_WINDOW_TOKENS,
497            max_output_tokens: COMPATIBILITY_MAX_OUTPUT_TOKENS,
498            safety_margin_tokens: COMPATIBILITY_SAFETY_MARGIN_TOKENS,
499            safe_window_percent: COMPATIBILITY_SAFE_WINDOW_PERCENT,
500            target_summary_tokens: previous_summary_tokens
501                .saturating_add(((represented_source_tokens as f64) * 0.20).ceil().max(1.0) as u32)
502                .max(1),
503            target_ratio: 0.20,
504        }
505    }
506
507    fn request_fits(
508        &self,
509        messages: &[Message],
510        requested_output_tokens: u32,
511        budget: &SummaryRequestBudget,
512    ) -> bool {
513        if requested_output_tokens == 0 || requested_output_tokens > budget.max_output_tokens {
514            return false;
515        }
516        let counter = TiktokenTokenCounter::default();
517        let input_tokens = counter.count_messages(messages);
518        input_tokens
519            .saturating_add(requested_output_tokens)
520            .saturating_add(budget.safety_margin_tokens)
521            <= budget.safe_request_tokens()
522    }
523
524    fn source_units(&self, messages: &[Message]) -> Vec<SourceUnit> {
525        let counter = TiktokenTokenCounter::default();
526        MessageSegmenter::new()
527            .segment(messages.to_vec())
528            .into_iter()
529            .filter(|segment| !segment.messages.is_empty())
530            .map(|segment| {
531                let first_message_id = segment
532                    .messages
533                    .first()
534                    .map(|message| message.id.clone())
535                    .unwrap_or_default();
536                let last_message_id = segment
537                    .messages
538                    .last()
539                    .map(|message| message.id.clone())
540                    .unwrap_or_else(|| first_message_id.clone());
541                SourceUnit {
542                    text: Self::render_messages(&segment.messages),
543                    represented_source_tokens: counter.count_messages(&segment.messages),
544                    first_message_id,
545                    last_message_id,
546                    continuation_part: 1,
547                }
548            })
549            .collect()
550    }
551
552    fn split_oversized_source_unit(
553        &self,
554        unit: SourceUnit,
555        budget: &SummaryRequestBudget,
556    ) -> Result<Vec<SourceUnit>, bamboo_compression::types::BudgetError> {
557        let counter = TiktokenTokenCounter::default();
558        let empty_prompt_tokens = counter.count_messages(&self.build_map_messages(&[], 1));
559        let available = budget
560            .safe_request_tokens()
561            .saturating_sub(budget.safety_margin_tokens)
562            .saturating_sub(empty_prompt_tokens);
563        if available < 2 {
564            return Err(bamboo_compression::types::BudgetError::TokenCountError(
565                "summarization model context is too small for map prompt overhead".to_string(),
566            ));
567        }
568
569        let ratio = budget.target_ratio.max(0.01);
570        let max_by_window = ((available as f64) / (1.0 + ratio)).floor() as u32;
571        let max_by_output = ((budget.max_output_tokens as f64) / ratio).floor() as u32;
572        let initial_piece_tokens = max_by_window.min(max_by_output).max(1);
573
574        let mut remaining = unit.text;
575        let mut remaining_represented = unit.represented_source_tokens;
576        let mut continuation_part = unit.continuation_part;
577        let mut parts = Vec::new();
578        while !remaining.is_empty() {
579            let remaining_text_tokens = counter.count_text(&remaining).max(1);
580            let mut piece_token_budget = initial_piece_tokens.min(remaining_text_tokens).max(1);
581            let (piece_text, piece_represented) = loop {
582                let mut prefix = counter.truncate_to_token_prefix(&remaining, piece_token_budget);
583                if prefix.is_empty() {
584                    prefix = remaining
585                        .chars()
586                        .next()
587                        .map(|character| character.to_string())
588                        .unwrap_or_default();
589                }
590                let is_last = prefix.len() == remaining.len();
591                let prefix_tokens = counter.count_text(&prefix).max(1);
592                let represented = if is_last {
593                    remaining_represented
594                } else if remaining_represented <= 1 {
595                    0
596                } else {
597                    (((remaining_represented as u64) * (prefix_tokens as u64)
598                        / (remaining_text_tokens as u64))
599                        .max(1)
600                        .min(remaining_represented.saturating_sub(1) as u64))
601                        as u32
602                };
603                let candidate = SourceUnit {
604                    text: prefix.clone(),
605                    represented_source_tokens: represented,
606                    first_message_id: unit.first_message_id.clone(),
607                    last_message_id: unit.last_message_id.clone(),
608                    continuation_part,
609                };
610                let requested_output = self.target_for_source(represented);
611                if self.request_fits(
612                    &self.build_map_messages(std::slice::from_ref(&candidate), requested_output),
613                    requested_output,
614                    budget,
615                ) {
616                    break (prefix, represented);
617                }
618                if piece_token_budget <= 1 {
619                    return Err(bamboo_compression::types::BudgetError::TokenCountError(
620                        format!(
621                            "single source continuation cannot fit summarization model window (range {}..{})",
622                            unit.first_message_id, unit.last_message_id
623                        ),
624                    ));
625                }
626                piece_token_budget = (piece_token_budget * 3 / 4).max(1);
627            };
628
629            let consumed_bytes = piece_text.len();
630            parts.push(SourceUnit {
631                text: piece_text,
632                represented_source_tokens: piece_represented,
633                first_message_id: unit.first_message_id.clone(),
634                last_message_id: unit.last_message_id.clone(),
635                continuation_part,
636            });
637            remaining = remaining[consumed_bytes..].to_string();
638            remaining_represented = remaining_represented.saturating_sub(piece_represented);
639            continuation_part += 1;
640        }
641        Ok(parts)
642    }
643
644    fn pack_source_chunks(
645        &self,
646        messages: &[Message],
647        budget: &SummaryRequestBudget,
648    ) -> Result<Vec<Vec<SourceUnit>>, bamboo_compression::types::BudgetError> {
649        let mut bounded_units = Vec::new();
650        for unit in self.source_units(messages) {
651            let requested_output = self.target_for_source(unit.represented_source_tokens);
652            let prompt = self.build_map_messages(std::slice::from_ref(&unit), requested_output);
653            if self.request_fits(&prompt, requested_output, budget) {
654                bounded_units.push(unit);
655            } else {
656                bounded_units.extend(self.split_oversized_source_unit(unit, budget)?);
657            }
658        }
659
660        let mut chunks = Vec::new();
661        let mut current = Vec::<SourceUnit>::new();
662        for unit in bounded_units {
663            let mut candidate = current.clone();
664            candidate.push(unit.clone());
665            let represented = candidate.iter().fold(0u32, |total, item| {
666                total.saturating_add(item.represented_source_tokens)
667            });
668            let requested_output = self.target_for_source(represented);
669            if self.request_fits(
670                &self.build_map_messages(&candidate, requested_output),
671                requested_output,
672                budget,
673            ) {
674                current = candidate;
675                continue;
676            }
677            if !current.is_empty() {
678                chunks.push(std::mem::take(&mut current));
679            }
680            let requested_output = self.target_for_source(unit.represented_source_tokens);
681            if !self.request_fits(
682                &self.build_map_messages(std::slice::from_ref(&unit), requested_output),
683                requested_output,
684                budget,
685            ) {
686                return Err(bamboo_compression::types::BudgetError::TokenCountError(
687                    "split source unit still exceeds summarization request ceiling".to_string(),
688                ));
689            }
690            current.push(unit);
691        }
692        if !current.is_empty() {
693            chunks.push(current);
694        }
695        if chunks.is_empty() {
696            return Err(bamboo_compression::types::BudgetError::TokenCountError(
697                "no bounded summarization chunks were produced".to_string(),
698            ));
699        }
700        Ok(chunks)
701    }
702
703    fn pack_reduction_groups(
704        &self,
705        parts: &[SummaryPart],
706        budget: &SummaryRequestBudget,
707    ) -> Vec<Vec<SummaryPart>> {
708        let mut groups = Vec::new();
709        let mut current = Vec::<SummaryPart>::new();
710        for part in parts {
711            let mut candidate = current.clone();
712            candidate.push(part.clone());
713            let represented = candidate.iter().fold(0u32, |total, item| {
714                total.saturating_add(item.represented_source_tokens)
715            });
716            let requested_output = self.target_for_source(represented);
717            if self.request_fits(
718                &self.build_reduce_messages(&candidate, requested_output, false),
719                requested_output,
720                budget,
721            ) {
722                current = candidate;
723            } else {
724                if !current.is_empty() {
725                    groups.push(std::mem::take(&mut current));
726                }
727                current.push(part.clone());
728            }
729        }
730        if !current.is_empty() {
731            groups.push(current);
732        }
733        groups
734    }
735
736    fn pack_multipart_final_groups(
737        &self,
738        parts: &[SummaryPart],
739        shared_context_capsule: &str,
740        retain_shared_context_capsule: bool,
741        budget: &SummaryRequestBudget,
742    ) -> Result<Vec<Vec<SummaryPart>>, bamboo_compression::types::BudgetError> {
743        let mut bounded_parts = Vec::new();
744        for part in parts {
745            let requested_output = self.target_for_source(part.represented_source_tokens);
746            if self.request_fits(
747                &self.build_multipart_finalize_messages(
748                    std::slice::from_ref(part),
749                    requested_output,
750                    shared_context_capsule,
751                    retain_shared_context_capsule,
752                ),
753                requested_output,
754                budget,
755            ) {
756                bounded_parts.push(part.clone());
757            } else {
758                bounded_parts.extend(self.split_oversized_multipart_part(
759                    part.clone(),
760                    shared_context_capsule,
761                    retain_shared_context_capsule,
762                    budget,
763                )?);
764            }
765        }
766
767        let mut groups = Vec::new();
768        let mut current = Vec::<SummaryPart>::new();
769        for part in &bounded_parts {
770            let mut candidate = current.clone();
771            candidate.push(part.clone());
772            let represented = candidate.iter().fold(0u32, |total, item| {
773                total.saturating_add(item.represented_source_tokens)
774            });
775            let requested_output = self.target_for_source(represented);
776            if self.request_fits(
777                &self.build_multipart_finalize_messages(
778                    &candidate,
779                    requested_output,
780                    shared_context_capsule,
781                    retain_shared_context_capsule,
782                ),
783                requested_output,
784                budget,
785            ) {
786                current = candidate;
787                continue;
788            }
789
790            if !current.is_empty() {
791                groups.push(std::mem::take(&mut current));
792            }
793            current.push(part.clone());
794        }
795        if !current.is_empty() {
796            groups.push(current);
797        }
798        Ok(groups)
799    }
800
801    fn split_oversized_multipart_part(
802        &self,
803        part: SummaryPart,
804        shared_context_capsule: &str,
805        retain_shared_context_capsule: bool,
806        budget: &SummaryRequestBudget,
807    ) -> Result<Vec<SummaryPart>, bamboo_compression::types::BudgetError> {
808        let counter = TiktokenTokenCounter::default();
809        let mut remaining = part.content;
810        let mut remaining_represented = part.represented_source_tokens;
811        let mut split_parts = Vec::new();
812        while !remaining.is_empty() {
813            let remaining_text_tokens = counter.count_text(&remaining).max(1);
814            let mut piece_token_budget = remaining_text_tokens;
815            let (piece_text, piece_represented) = loop {
816                let mut prefix = counter.truncate_to_token_prefix(&remaining, piece_token_budget);
817                if prefix.is_empty() {
818                    prefix = remaining
819                        .chars()
820                        .next()
821                        .map(|character| character.to_string())
822                        .unwrap_or_default();
823                }
824                let is_last = prefix.len() == remaining.len();
825                let prefix_tokens = counter.count_text(&prefix).max(1);
826                let represented = if is_last {
827                    remaining_represented
828                } else if remaining_represented <= 1 {
829                    0
830                } else {
831                    (((remaining_represented as u64) * (prefix_tokens as u64)
832                        / (remaining_text_tokens as u64))
833                        .max(1)
834                        .min(remaining_represented.saturating_sub(1) as u64))
835                        as u32
836                };
837                let candidate = SummaryPart {
838                    content: prefix.clone(),
839                    represented_source_tokens: represented,
840                    first_message_id: part.first_message_id.clone(),
841                    last_message_id: part.last_message_id.clone(),
842                };
843                let requested_output = self.target_for_source(represented);
844                if self.request_fits(
845                    &self.build_multipart_finalize_messages(
846                        std::slice::from_ref(&candidate),
847                        requested_output,
848                        shared_context_capsule,
849                        retain_shared_context_capsule,
850                    ),
851                    requested_output,
852                    budget,
853                ) {
854                    break (prefix, represented);
855                }
856                if piece_token_budget <= 1 {
857                    return Err(bamboo_compression::types::BudgetError::TokenCountError(
858                        format!(
859                            "shared-context capsule and multipart source range {}..{} cannot fit the summarization request ceiling",
860                            part.first_message_id, part.last_message_id
861                        ),
862                    ));
863                }
864                piece_token_budget = (piece_token_budget * 3 / 4).max(1);
865            };
866
867            let consumed_bytes = piece_text.len();
868            split_parts.push(SummaryPart {
869                content: piece_text,
870                represented_source_tokens: piece_represented,
871                first_message_id: part.first_message_id.clone(),
872                last_message_id: part.last_message_id.clone(),
873            });
874            remaining = remaining[consumed_bytes..].to_string();
875            remaining_represented = remaining_represented.saturating_sub(piece_represented);
876        }
877        Ok(split_parts)
878    }
879
880    async fn execute_bounded_request(
881        &self,
882        stage: &str,
883        stage_index: usize,
884        stage_count: usize,
885        messages: &[Message],
886        requested_output_tokens: u32,
887        budget: &SummaryRequestBudget,
888    ) -> Result<String, bamboo_compression::types::BudgetError> {
889        let counter = TiktokenTokenCounter::default();
890        let input_tokens = counter.count_messages(messages);
891        let safe_request_tokens = budget.safe_request_tokens();
892        if !self.request_fits(messages, requested_output_tokens, budget) {
893            return Err(bamboo_compression::types::BudgetError::TokenCountError(
894                format!(
895                    "bounded summarization invariant failed for {stage}: input={input_tokens}, output={requested_output_tokens}, safety={}, safe_limit={safe_request_tokens}",
896                    budget.safety_margin_tokens,
897                ),
898            ));
899        }
900        let progress = SummarizationProgress {
901            stage: stage.to_string(),
902            stage_index,
903            stage_count,
904            estimated_input_tokens: input_tokens,
905            requested_output_tokens,
906            safe_request_tokens,
907            model_context_tokens: budget.context_window_tokens,
908        };
909        if let Some(callback) = self.progress_callback.as_ref() {
910            callback(&progress);
911        }
912        tracing::info!(
913            logical_pass_id = self.logical_pass_id.as_deref().unwrap_or("untracked"),
914            logical_phase = self.logical_phase.as_deref().unwrap_or("unspecified"),
915            stage,
916            stage_index,
917            stage_count,
918            input_tokens,
919            requested_output_tokens,
920            safe_request_tokens,
921            model_context_tokens = budget.context_window_tokens,
922            model = %self.model,
923            "Executing bounded summarization request"
924        );
925        let content = self
926            .collect_stream_response(messages, requested_output_tokens)
927            .await
928            .map_err(|error| {
929                bamboo_compression::types::BudgetError::TokenCountError(format!(
930                    "{stage} stage {stage_index}/{stage_count} failed: {error}"
931                ))
932            })?;
933        if content.trim().is_empty() {
934            return Err(bamboo_compression::types::BudgetError::TokenCountError(
935                format!(
936                    "{stage} stage {stage_index}/{stage_count} returned an empty completed response"
937                ),
938            ));
939        }
940        Ok(content)
941    }
942
943    async fn build_bounded_shared_context_capsule(
944        &self,
945        shared_context: &str,
946        target_tokens: u32,
947        budget: &SummaryRequestBudget,
948    ) -> Result<(String, u32, u32), bamboo_compression::types::BudgetError> {
949        let counter = TiktokenTokenCounter::default();
950        let shared_tokens = counter.count_text(shared_context).max(1);
951        let capsule_target = target_tokens.max(1).min(shared_tokens);
952        // A shared-context capsule is still a compression pass, so it follows
953        // the same map-then-reduce contract even when its fully rendered source
954        // would fit one provider request. The child has no shared context of its
955        // own, so multipart finalization cannot recurse back into this lane.
956        let capsule_messages = [Message::user(format!(
957            "Preserve this shared context capsule, including every binding compression \
958             instruction and hook-injected directive:\n\n{shared_context}"
959        ))];
960        let capsule_source_tokens = counter.count_messages(&capsule_messages).max(1);
961        let mut capsule_budget = budget.clone();
962        capsule_budget.target_summary_tokens = capsule_target;
963        // A prior durable summary has already paid the global compression
964        // ratio. Allow a 1:1 capsule map when its retained allocation calls for
965        // it; clamping this child to the normal 50% ceiling would silently apply
966        // a second compression ratio before the outer multipart reduce.
967        capsule_budget.target_ratio =
968            (capsule_target as f64 / capsule_source_tokens as f64).min(1.0);
969        let mut capsule_summarizer =
970            LlmSummarizer::new(Arc::clone(&self.llm), self.model.clone(), None, None)
971                .with_summary_mode(SummaryMode::FullRewrite)
972                .with_request_budget(capsule_budget)
973                .with_heuristic_fallback_on_error(false);
974        if let (Some(pass_id), Some(phase)) =
975            (self.logical_pass_id.as_ref(), self.logical_phase.as_ref())
976        {
977            capsule_summarizer =
978                capsule_summarizer.with_logical_pass_context(pass_id.clone(), phase.clone());
979        }
980        if let Some(callback) = self.progress_callback.as_ref() {
981            capsule_summarizer = capsule_summarizer.with_progress_callback(Arc::clone(callback));
982        }
983        let report = Box::pin(capsule_summarizer.summarize_with_report(&capsule_messages)).await?;
984        let content = counter.truncate_to_token_prefix(report.content.trim(), capsule_target);
985        if content.trim().is_empty() {
986            return Err(bamboo_compression::types::BudgetError::TokenCountError(
987                "bounded shared-context capsule is empty".to_string(),
988            ));
989        }
990        Ok((content, report.map_calls, report.reduce_calls))
991    }
992
993    fn compose_multipart_summary(
994        parts: &[SummaryPart],
995        shared_context_capsule: Option<&str>,
996    ) -> String {
997        let mut sections = Vec::new();
998        if let Some(capsule) = shared_context_capsule
999            .map(str::trim)
1000            .filter(|capsule| !capsule.is_empty())
1001        {
1002            sections.push(capsule.to_string());
1003        }
1004        sections.extend(
1005            parts
1006                .iter()
1007                .map(|part| part.content.trim())
1008                .filter(|content| !content.is_empty())
1009                .map(String::from),
1010        );
1011        sections.join("\n\n")
1012    }
1013
1014    async fn finalize_multipart_summary(
1015        &self,
1016        parts: &[SummaryPart],
1017        budget: &SummaryRequestBudget,
1018    ) -> Result<(String, u32, u32), bamboo_compression::types::BudgetError> {
1019        let shared_context = self.render_shared_context();
1020        let represented_new_source = parts.iter().fold(0u32, |total, part| {
1021            total.saturating_add(part.represented_source_tokens)
1022        });
1023        let new_source_target = self.target_for_source(represented_new_source);
1024        let retained_capsule_target = budget
1025            .target_summary_tokens
1026            .saturating_sub(new_source_target);
1027        let retain_shared_context_capsule = retained_capsule_target > 0;
1028        let (shared_context_capsule, capsule_map_calls, capsule_reduce_calls) =
1029            if shared_context.trim().is_empty() {
1030                (String::new(), 0, 0)
1031            } else {
1032                let counter = TiktokenTokenCounter::default();
1033                let reference_only_target =
1034                    self.target_for_source(counter.count_text(&shared_context));
1035                self.build_bounded_shared_context_capsule(
1036                    &shared_context,
1037                    if retain_shared_context_capsule {
1038                        retained_capsule_target
1039                    } else {
1040                        reference_only_target
1041                    },
1042                    budget,
1043                )
1044                .await?
1045            };
1046        let mut multipart_reduce_calls = 0u32;
1047        let groups = self.pack_multipart_final_groups(
1048            parts,
1049            &shared_context_capsule,
1050            retain_shared_context_capsule,
1051            budget,
1052        )?;
1053        let group_count = groups.len();
1054        let mut finalized = Vec::with_capacity(group_count);
1055        for (group_index, group) in groups.into_iter().enumerate() {
1056            let represented = group.iter().fold(0u32, |total, part| {
1057                total.saturating_add(part.represented_source_tokens)
1058            });
1059            let requested_output = self.target_for_source(represented);
1060            let prompt = self.build_multipart_finalize_messages(
1061                &group,
1062                requested_output,
1063                &shared_context_capsule,
1064                retain_shared_context_capsule,
1065            );
1066            let content = self
1067                .execute_bounded_request(
1068                    "multipart_final_reduce",
1069                    group_index + 1,
1070                    group_count,
1071                    &prompt,
1072                    requested_output,
1073                    budget,
1074                )
1075                .await?;
1076            multipart_reduce_calls = multipart_reduce_calls.saturating_add(1);
1077            finalized.push(SummaryPart {
1078                content,
1079                represented_source_tokens: represented,
1080                first_message_id: group
1081                    .first()
1082                    .map(|part| part.first_message_id.clone())
1083                    .unwrap_or_default(),
1084                last_message_id: group
1085                    .last()
1086                    .map(|part| part.last_message_id.clone())
1087                    .unwrap_or_default(),
1088            });
1089        }
1090        Ok((
1091            Self::compose_multipart_summary(
1092                &finalized,
1093                (retain_shared_context_capsule && !shared_context_capsule.trim().is_empty())
1094                    .then_some(shared_context_capsule.as_str()),
1095            ),
1096            capsule_map_calls,
1097            capsule_reduce_calls.saturating_add(multipart_reduce_calls),
1098        ))
1099    }
1100
1101    async fn summarize_bounded(
1102        &self,
1103        messages: &[Message],
1104        budget: &SummaryRequestBudget,
1105    ) -> Result<SummarizationReport, bamboo_compression::types::BudgetError> {
1106        let counter = TiktokenTokenCounter::default();
1107        let represented_source_tokens = counter.count_messages(messages);
1108        let final_target = budget.target_summary_tokens.max(1);
1109
1110        // Every bounded compression pass deliberately goes through map then
1111        // reduce, even when the source would fit a single provider request.
1112        // Keeping one pipeline for all source sizes avoids a size-dependent
1113        // semantic split and guarantees that no raw source is ever sent
1114        // directly to the terminal summarizer.
1115        let chunks = self.pack_source_chunks(messages, budget)?;
1116        let mut map_calls = 0u32;
1117        let mut parts = Vec::with_capacity(chunks.len());
1118        for (index, chunk) in chunks.iter().enumerate() {
1119            let represented = chunk.iter().fold(0u32, |total, unit| {
1120                total.saturating_add(unit.represented_source_tokens)
1121            });
1122            let requested_output = self.target_for_source(represented);
1123            let prompt = self.build_map_messages(chunk, requested_output);
1124            let content = self
1125                .execute_bounded_request(
1126                    "map",
1127                    index + 1,
1128                    chunks.len(),
1129                    &prompt,
1130                    requested_output,
1131                    budget,
1132                )
1133                .await?;
1134            map_calls = map_calls.saturating_add(1);
1135            let first_message_id = chunk
1136                .first()
1137                .map(|unit| unit.first_message_id.clone())
1138                .unwrap_or_else(|| format!("chunk-{index}"));
1139            let last_message_id = chunk
1140                .last()
1141                .map(|unit| unit.last_message_id.clone())
1142                .unwrap_or_else(|| first_message_id.clone());
1143            parts.push(SummaryPart {
1144                content,
1145                represented_source_tokens: represented,
1146                first_message_id,
1147                last_message_id,
1148            });
1149        }
1150
1151        let mut reduce_calls = 0u32;
1152        let mut depth = 0usize;
1153        loop {
1154            let final_prompt = self.build_reduce_messages(&parts, final_target, true);
1155            if self.request_fits(&final_prompt, final_target, budget) {
1156                let content = self
1157                    .execute_bounded_request(
1158                        "final_reduce",
1159                        1,
1160                        1,
1161                        &final_prompt,
1162                        final_target,
1163                        budget,
1164                    )
1165                    .await?;
1166                reduce_calls = reduce_calls.saturating_add(1);
1167                let actual_summary_tokens = counter.count_text(&content);
1168                let underfilled =
1169                    actual_summary_tokens.saturating_mul(5) < final_target.saturating_mul(4);
1170                return Ok(SummarizationReport {
1171                    content,
1172                    represented_source_tokens,
1173                    target_summary_tokens: final_target,
1174                    actual_summary_tokens,
1175                    map_calls,
1176                    reduce_calls,
1177                    fallback_used: false,
1178                    budget_clamped: underfilled,
1179                    budget_clamp_reason: underfilled
1180                        .then(|| "model_returned_below_80_percent_of_target".to_string()),
1181                });
1182            }
1183
1184            if parts.len() <= 1 || depth >= 8 {
1185                break;
1186            }
1187            let groups = self.pack_reduction_groups(&parts, budget);
1188            if groups.len() >= parts.len() {
1189                break;
1190            }
1191            let mut reduced = Vec::with_capacity(groups.len());
1192            let group_count = groups.len();
1193            for (group_index, group) in groups.into_iter().enumerate() {
1194                if group.len() == 1 {
1195                    reduced.push(group.into_iter().next().expect("single reduction part"));
1196                    continue;
1197                }
1198                let represented = group.iter().fold(0u32, |total, part| {
1199                    total.saturating_add(part.represented_source_tokens)
1200                });
1201                let requested_output = self.target_for_source(represented);
1202                let prompt = self.build_reduce_messages(&group, requested_output, false);
1203                let content = self
1204                    .execute_bounded_request(
1205                        "intermediate_reduce",
1206                        group_index + 1,
1207                        group_count,
1208                        &prompt,
1209                        requested_output,
1210                        budget,
1211                    )
1212                    .await?;
1213                reduce_calls = reduce_calls.saturating_add(1);
1214                reduced.push(SummaryPart {
1215                    content,
1216                    represented_source_tokens: represented,
1217                    first_message_id: group
1218                        .first()
1219                        .map(|part| part.first_message_id.clone())
1220                        .unwrap_or_default(),
1221                    last_message_id: group
1222                        .last()
1223                        .map(|part| part.last_message_id.clone())
1224                        .unwrap_or_default(),
1225                });
1226            }
1227            parts = reduced;
1228            depth += 1;
1229        }
1230
1231        // A single final model response cannot represent an arbitrarily large
1232        // 20%-of-source target when the selected model has a smaller output
1233        // limit. Finalize bounded sections under one shared-context capsule,
1234        // then compose them instead of applying another global 20% pass (which
1235        // would collapse to ~4%).
1236        let (content, multipart_map_calls, multipart_reduce_calls) =
1237            self.finalize_multipart_summary(&parts, budget).await?;
1238        map_calls = map_calls.saturating_add(multipart_map_calls);
1239        reduce_calls = reduce_calls.saturating_add(multipart_reduce_calls);
1240        let actual_summary_tokens = counter.count_text(&content);
1241        let underfilled = actual_summary_tokens.saturating_mul(5) < final_target.saturating_mul(4);
1242        tracing::info!(
1243            logical_pass_id = self.logical_pass_id.as_deref().unwrap_or("untracked"),
1244            logical_phase = self.logical_phase.as_deref().unwrap_or("unspecified"),
1245            part_count = parts.len(),
1246            multipart_map_calls,
1247            multipart_reduce_calls,
1248            actual_summary_tokens,
1249            target_summary_tokens = final_target,
1250            "Final single-response reduce did not fit; composed instruction-aware bounded multipart summary"
1251        );
1252        Ok(SummarizationReport {
1253            content,
1254            represented_source_tokens,
1255            target_summary_tokens: final_target,
1256            actual_summary_tokens,
1257            map_calls,
1258            reduce_calls,
1259            fallback_used: false,
1260            budget_clamped: underfilled,
1261            budget_clamp_reason: underfilled
1262                .then(|| "multipart_summary_below_80_percent_of_target".to_string()),
1263        })
1264    }
1265
1266    async fn heuristic_report(
1267        &self,
1268        messages: &[Message],
1269        target_summary_tokens: u32,
1270        reason: &str,
1271    ) -> Result<SummarizationReport, bamboo_compression::types::BudgetError> {
1272        let counter = TiktokenTokenCounter::default();
1273        let heuristic = HeuristicSummarizer::new().summarize(messages).await?;
1274        let content = counter.truncate_to_token_prefix(&heuristic, target_summary_tokens.max(1));
1275        Ok(SummarizationReport {
1276            represented_source_tokens: counter.count_messages(messages),
1277            actual_summary_tokens: counter.count_text(&content),
1278            target_summary_tokens,
1279            content,
1280            map_calls: 0,
1281            reduce_calls: 0,
1282            fallback_used: true,
1283            budget_clamped: true,
1284            budget_clamp_reason: Some(reason.to_string()),
1285        })
1286    }
1287
1288    pub async fn summarize_with_report(
1289        &self,
1290        messages: &[Message],
1291    ) -> Result<SummarizationReport, bamboo_compression::types::BudgetError> {
1292        if messages.is_empty() {
1293            return Ok(SummarizationReport {
1294                content: "No conversation history to summarize.".to_string(),
1295                represented_source_tokens: 0,
1296                target_summary_tokens: 0,
1297                actual_summary_tokens: 0,
1298                map_calls: 0,
1299                reduce_calls: 0,
1300                fallback_used: false,
1301                budget_clamped: false,
1302                budget_clamp_reason: None,
1303            });
1304        }
1305
1306        let compatibility_budget;
1307        let budget = if let Some(budget) = self.request_budget.as_ref() {
1308            budget
1309        } else {
1310            compatibility_budget = self.compatibility_request_budget(messages);
1311            &compatibility_budget
1312        };
1313        let target_summary_tokens = budget.target_summary_tokens;
1314        let result = self.summarize_bounded(messages, budget).await;
1315
1316        match result {
1317            Ok(report) if !report.content.trim().is_empty() => Ok(report),
1318            Ok(_) => {
1319                tracing::warn!(
1320                    "LlmSummarizer: LLM returned empty summary, falling back to heuristic"
1321                );
1322                self.heuristic_report(messages, target_summary_tokens, "empty_llm_response")
1323                    .await
1324            }
1325            Err(error) if self.heuristic_fallback_on_error && self.request_budget.is_none() => {
1326                tracing::warn!(
1327                    "LlmSummarizer: compatibility map/reduce pipeline failed ({}), falling back to heuristic",
1328                    error
1329                );
1330                self.heuristic_report(
1331                    messages,
1332                    target_summary_tokens,
1333                    "llm_error_heuristic_fallback",
1334                )
1335                .await
1336            }
1337            Err(error) => Err(error),
1338        }
1339    }
1340
1341    /// Consume an LLM stream and collect the full text response.
1342    async fn collect_stream_response(
1343        &self,
1344        messages: &[Message],
1345        max_output_tokens: u32,
1346    ) -> Result<String, bamboo_compression::types::BudgetError> {
1347        // Compression calls need most of their output allowance for the summary
1348        // itself. Low reasoning keeps dynamic, potentially small output budgets
1349        // useful across both reasoning and non-reasoning models.
1350        let options = LLMRequestOptions {
1351            session_id: None,
1352            reasoning_effort: Some(ReasoningEffort::Low),
1353            parallel_tool_calls: None,
1354            required_tool: None,
1355            responses: None,
1356            request_purpose: Some("compression".to_string()),
1357            cache: None,
1358        };
1359        let stream = self
1360            .llm
1361            .chat_stream_with_options(
1362                messages,
1363                &[],
1364                Some(max_output_tokens.max(1)),
1365                &self.model,
1366                Some(&options),
1367            )
1368            .await
1369            .map_err(|e| {
1370                bamboo_compression::types::BudgetError::TokenCountError(format!(
1371                    "LLM summarization call failed: {}",
1372                    e
1373                ))
1374            })?;
1375
1376        let mut content = String::new();
1377        let mut stream = stream;
1378        let mut terminal_done = false;
1379
1380        while let Some(chunk_result) = stream.next().await {
1381            match chunk_result {
1382                Ok(LLMChunk::Token(text)) => content.push_str(&text),
1383                Ok(LLMChunk::Done) => {
1384                    terminal_done = true;
1385                    break;
1386                }
1387                Ok(_) => {} // Ignore reasoning tokens, tool calls, etc.
1388                Err(e) => {
1389                    tracing::warn!("LLM summarization stream error: {}", e);
1390                    return Err(bamboo_compression::types::BudgetError::TokenCountError(
1391                        format!("LLM summarization stream failed: {}", e),
1392                    ));
1393                }
1394            }
1395        }
1396
1397        if !terminal_done && !content.is_empty() {
1398            return Err(bamboo_compression::types::BudgetError::TokenCountError(
1399                "LLM summarization stream ended without terminal completion".to_string(),
1400            ));
1401        }
1402        Ok(content)
1403    }
1404}
1405
1406impl std::fmt::Debug for LlmSummarizer {
1407    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1408        f.debug_struct("LlmSummarizer")
1409            .field("model", &self.model)
1410            .field("has_existing_summary", &self.existing_summary.is_some())
1411            .field("context_block_count", &self.context_blocks.len())
1412            .field("logical_pass_id", &self.logical_pass_id)
1413            .field("logical_phase", &self.logical_phase)
1414            .field("has_progress_callback", &self.progress_callback.is_some())
1415            .finish()
1416    }
1417}
1418
1419#[async_trait]
1420impl Summarizer for LlmSummarizer {
1421    async fn summarize(
1422        &self,
1423        messages: &[Message],
1424    ) -> Result<String, bamboo_compression::types::BudgetError> {
1425        tracing::info!(
1426            "LlmSummarizer: summarizing {} messages using model '{}' (existing_summary={})",
1427            messages.len(),
1428            self.model,
1429            self.existing_summary.is_some()
1430        );
1431        let report = self.summarize_with_report(messages).await?;
1432        tracing::info!(
1433            chars = report.content.len(),
1434            actual_summary_tokens = report.actual_summary_tokens,
1435            target_summary_tokens = report.target_summary_tokens,
1436            map_calls = report.map_calls,
1437            reduce_calls = report.reduce_calls,
1438            fallback_used = report.fallback_used,
1439            "LlmSummarizer: generated summary"
1440        );
1441        Ok(report.content)
1442    }
1443
1444    fn estimate_summary_tokens(&self, message_count: usize) -> u32 {
1445        self.request_budget
1446            .as_ref()
1447            .map(|budget| budget.target_summary_tokens)
1448            .unwrap_or_else(|| (message_count * 80).min(2000) as u32)
1449    }
1450}
1451
1452#[cfg(test)]
1453mod tests {
1454    use super::*;
1455    use bamboo_domain::{FunctionCall, ReasoningEffort, ToolCall};
1456    use bamboo_llm::{LLMChunk, LLMError, LLMRequestOptions, LLMStream};
1457    use futures::stream;
1458    use std::sync::Mutex;
1459
1460    struct DummyProvider;
1461
1462    #[async_trait]
1463    impl LLMProvider for DummyProvider {
1464        async fn chat_stream(
1465            &self,
1466            _messages: &[Message],
1467            _tools: &[bamboo_domain::ToolSchema],
1468            _max_output_tokens: Option<u32>,
1469            _model: &str,
1470        ) -> Result<LLMStream, LLMError> {
1471            Ok(Box::pin(stream::iter(vec![
1472                Ok::<LLMChunk, LLMError>(LLMChunk::Token("dummy summary".to_string())),
1473                Ok::<LLMChunk, LLMError>(LLMChunk::Done),
1474            ])))
1475        }
1476    }
1477
1478    fn test_summary_part() -> SummaryPart {
1479        SummaryPart {
1480            content: "New chronological partial summary".to_string(),
1481            represented_source_tokens: 100,
1482            first_message_id: "first".to_string(),
1483            last_message_id: "last".to_string(),
1484        }
1485    }
1486
1487    #[test]
1488    fn llm_summarizer_prompt_includes_context_blocks_and_state_sections() {
1489        let summarizer = LlmSummarizer::new(
1490            Arc::new(DummyProvider),
1491            "gpt-4o-mini".to_string(),
1492            Some("Earlier summary".to_string()),
1493            Some(
1494                "## Current Task List\n[/] task_1: Fix compression bounce\n[x] task_0: Analyze bug"
1495                    .to_string(),
1496            ),
1497        )
1498        .with_context_blocks(vec![
1499            ContextBlock::new(
1500                ContextBlockType::TaskSnapshot,
1501                ContextBlockPriority::High,
1502                ContextBlockStability::RoundDynamic,
1503                "Current Task List",
1504                "[/] task_1: Fix compression bounce",
1505            ),
1506            ContextBlock::new(
1507                ContextBlockType::ExternalMemory,
1508                ContextBlockPriority::Medium,
1509                ContextBlockStability::RoundDynamic,
1510                "External Memory (Persistent)",
1511                "Session note body",
1512            ),
1513        ]);
1514        let prompt_messages = summarizer.build_reduce_messages(&[test_summary_part()], 20, true);
1515        assert_eq!(prompt_messages.len(), 2);
1516        assert_eq!(prompt_messages[0].role, Role::System);
1517        assert!(prompt_messages[1]
1518            .content
1519            .contains("## Compression Context Blocks"));
1520        assert!(prompt_messages[1].content.contains("Current Task List"));
1521        assert!(prompt_messages[1]
1522            .content
1523            .contains("External Memory (Persistent)"));
1524        assert!(prompt_messages[1]
1525            .content
1526            .contains("Current active objective"));
1527        assert!(prompt_messages[1].content.contains("Requirement checklist"));
1528        assert!(prompt_messages[1].content.contains("Active tasks"));
1529        assert!(prompt_messages[1].content.contains("Completed tasks"));
1530        assert!(prompt_messages[1]
1531            .content
1532            .contains("Obsolete or superseded tasks"));
1533        assert!(prompt_messages[1].content.contains("Earlier summary"));
1534    }
1535
1536    #[derive(Default)]
1537    struct ReasoningCaptureProvider {
1538        captured_reasoning: Mutex<Vec<Option<ReasoningEffort>>>,
1539    }
1540
1541    #[async_trait]
1542    impl LLMProvider for ReasoningCaptureProvider {
1543        async fn chat_stream(
1544            &self,
1545            _messages: &[Message],
1546            _tools: &[bamboo_domain::ToolSchema],
1547            _max_output_tokens: Option<u32>,
1548            _model: &str,
1549        ) -> Result<LLMStream, LLMError> {
1550            Ok(Box::pin(stream::iter(vec![
1551                Ok::<LLMChunk, LLMError>(LLMChunk::Token("captured summary".to_string())),
1552                Ok::<LLMChunk, LLMError>(LLMChunk::Done),
1553            ])))
1554        }
1555
1556        async fn chat_stream_with_options(
1557            &self,
1558            messages: &[Message],
1559            tools: &[bamboo_domain::ToolSchema],
1560            max_output_tokens: Option<u32>,
1561            model: &str,
1562            options: Option<&LLMRequestOptions>,
1563        ) -> Result<LLMStream, LLMError> {
1564            self.captured_reasoning
1565                .lock()
1566                .expect("captured reasoning lock should not be poisoned")
1567                .push(options.and_then(|o| o.reasoning_effort));
1568            self.chat_stream(messages, tools, max_output_tokens, model)
1569                .await
1570        }
1571    }
1572
1573    #[tokio::test]
1574    async fn llm_summarizer_requests_low_reasoning_effort_for_summary_calls() {
1575        let provider = Arc::new(ReasoningCaptureProvider::default());
1576        let summarizer = LlmSummarizer::new(
1577            provider.clone(),
1578            "gpt-5-mini".to_string(),
1579            None,
1580            Some("task list".to_string()),
1581        );
1582        let messages = vec![
1583            Message::user("请总结最近三轮"),
1584            Message::assistant("已完成第一步并准备第二步", None),
1585        ];
1586
1587        let summary = summarizer
1588            .summarize(&messages)
1589            .await
1590            .expect("summary generation should succeed");
1591        assert_eq!(summary, "captured summary");
1592
1593        let captured = provider
1594            .captured_reasoning
1595            .lock()
1596            .expect("captured reasoning lock should not be poisoned");
1597        assert_eq!(
1598            captured.as_slice(),
1599            [Some(ReasoningEffort::Low), Some(ReasoningEffort::Low)],
1600            "compatibility callers must use one map request followed by one reduce request"
1601        );
1602    }
1603
1604    /// Provider that captures both `reasoning_effort` and `max_output_tokens`.
1605    #[derive(Default)]
1606    struct RequestOptionsCaptureProvider {
1607        captured_reasoning: Mutex<Vec<Option<ReasoningEffort>>>,
1608        captured_max_tokens: Mutex<Vec<Option<u32>>>,
1609    }
1610
1611    #[async_trait]
1612    impl LLMProvider for RequestOptionsCaptureProvider {
1613        async fn chat_stream(
1614            &self,
1615            _messages: &[Message],
1616            _tools: &[bamboo_domain::ToolSchema],
1617            _max_output_tokens: Option<u32>,
1618            _model: &str,
1619        ) -> Result<LLMStream, LLMError> {
1620            Ok(Box::pin(stream::iter(vec![
1621                Ok::<LLMChunk, LLMError>(LLMChunk::Token("captured summary".to_string())),
1622                Ok::<LLMChunk, LLMError>(LLMChunk::Done),
1623            ])))
1624        }
1625
1626        async fn chat_stream_with_options(
1627            &self,
1628            messages: &[Message],
1629            tools: &[bamboo_domain::ToolSchema],
1630            max_output_tokens: Option<u32>,
1631            model: &str,
1632            options: Option<&LLMRequestOptions>,
1633        ) -> Result<LLMStream, LLMError> {
1634            self.captured_reasoning
1635                .lock()
1636                .expect("lock should not be poisoned")
1637                .push(options.and_then(|o| o.reasoning_effort));
1638            self.captured_max_tokens
1639                .lock()
1640                .expect("lock should not be poisoned")
1641                .push(max_output_tokens);
1642            self.chat_stream(messages, tools, max_output_tokens, model)
1643                .await
1644        }
1645    }
1646
1647    #[tokio::test]
1648    async fn compatibility_summarizer_uses_source_derived_output_budget_with_low_reasoning() {
1649        let provider = Arc::new(RequestOptionsCaptureProvider::default());
1650        let summarizer = LlmSummarizer::new(
1651            provider.clone(),
1652            "gpt-5-mini".to_string(),
1653            None,
1654            Some("task list".to_string()),
1655        );
1656        let messages = vec![
1657            Message::user("请总结最近三轮"),
1658            Message::assistant("已完成第一步并准备第二步", None),
1659        ];
1660
1661        let summary = summarizer
1662            .summarize(&messages)
1663            .await
1664            .expect("summary generation should succeed");
1665        assert_eq!(summary, "captured summary");
1666
1667        let captured_reasoning = provider
1668            .captured_reasoning
1669            .lock()
1670            .expect("lock should not be poisoned");
1671        let captured_max_tokens = provider
1672            .captured_max_tokens
1673            .lock()
1674            .expect("lock should not be poisoned");
1675        assert_eq!(
1676            captured_reasoning.as_slice(),
1677            [Some(ReasoningEffort::Low), Some(ReasoningEffort::Low)]
1678        );
1679        let expected_target = ((TiktokenTokenCounter::default().count_messages(&messages) as f64)
1680            * 0.20)
1681            .ceil()
1682            .max(1.0) as u32;
1683        assert_eq!(
1684            captured_max_tokens.as_slice(),
1685            [Some(expected_target), Some(expected_target)],
1686            "both compatibility map and reduce requests use the source-derived 20% target"
1687        );
1688    }
1689
1690    #[test]
1691    fn compatibility_budget_retains_previous_summary_before_adding_twenty_percent() {
1692        let previous_summary = "durable prior summary evidence ".repeat(20);
1693        let summarizer = LlmSummarizer::new(
1694            Arc::new(DummyProvider),
1695            "summary-model".to_string(),
1696            Some(previous_summary.clone()),
1697            None,
1698        );
1699        let messages = summary_messages();
1700        let counter = TiktokenTokenCounter::default();
1701
1702        let budget = summarizer.compatibility_request_budget(&messages);
1703
1704        let expected = counter.count_text(&previous_summary).saturating_add(
1705            ((counter.count_messages(&messages) as f64) * 0.20)
1706                .ceil()
1707                .max(1.0) as u32,
1708        );
1709        assert_eq!(budget.target_summary_tokens, expected);
1710    }
1711
1712    #[test]
1713    fn full_rewrite_mode_uses_default_final_reduce_prompt() {
1714        let summarizer =
1715            LlmSummarizer::new(Arc::new(DummyProvider), "model".to_string(), None, None)
1716                .with_summary_mode(SummaryMode::FullRewrite);
1717        let prompts = summarizer.build_reduce_messages(&[test_summary_part()], 20, true);
1718        let system = &prompts[0].content;
1719        assert!(
1720            system.contains("final reduce stage"),
1721            "FullRewrite prompt should identify the final reduce stage"
1722        );
1723        assert!(
1724            !system.contains("incremental final reduce"),
1725            "FullRewrite prompt should not contain incremental language"
1726        );
1727    }
1728
1729    #[test]
1730    fn incremental_merge_mode_uses_update_final_reduce_prompt() {
1731        let summarizer = LlmSummarizer::new(
1732            Arc::new(DummyProvider),
1733            "model".to_string(),
1734            Some("Previous summary content".to_string()),
1735            None,
1736        )
1737        .with_summary_mode(SummaryMode::IncrementalMerge);
1738        let prompts = summarizer.build_reduce_messages(&[test_summary_part()], 20, true);
1739        let system = &prompts[0].content;
1740        assert!(
1741            system.contains("incremental final reduce stage"),
1742            "IncrementalMerge prompt should identify the incremental final reduce stage"
1743        );
1744        assert!(
1745            system.contains("Update the supplied prior summary"),
1746            "IncrementalMerge prompt should direct the reducer to update the prior summary"
1747        );
1748    }
1749
1750    #[test]
1751    fn default_summary_mode_is_full_rewrite() {
1752        assert!(matches!(SummaryMode::default(), SummaryMode::FullRewrite));
1753    }
1754
1755    #[test]
1756    fn incremental_merge_includes_existing_summary_in_user_content() {
1757        let summarizer = LlmSummarizer::new(
1758            Arc::new(DummyProvider),
1759            "model".to_string(),
1760            Some("Previous summary content".to_string()),
1761            None,
1762        )
1763        .with_summary_mode(SummaryMode::IncrementalMerge);
1764        let prompts = summarizer.build_reduce_messages(&[test_summary_part()], 20, true);
1765        let user_content = &prompts[1].content;
1766        assert!(
1767            user_content.contains("Previous Summary"),
1768            "IncrementalMerge user prompt should include the existing summary"
1769        );
1770        assert!(
1771            user_content.contains("Previous summary content"),
1772            "IncrementalMerge user prompt should include the actual summary text"
1773        );
1774    }
1775
1776    /// Provider whose summarization stream call fails transiently (500/429/timeout).
1777    struct FailingProvider;
1778
1779    #[async_trait]
1780    impl LLMProvider for FailingProvider {
1781        async fn chat_stream(
1782            &self,
1783            _messages: &[Message],
1784            _tools: &[bamboo_domain::ToolSchema],
1785            _max_output_tokens: Option<u32>,
1786            _model: &str,
1787        ) -> Result<LLMStream, LLMError> {
1788            Err(LLMError::Api("http 500 transient".to_string()))
1789        }
1790    }
1791
1792    fn summary_messages() -> Vec<Message> {
1793        vec![
1794            Message::user("do the work"),
1795            Message::assistant("working on it", None),
1796            Message::user("keep going"),
1797        ]
1798    }
1799
1800    #[tokio::test]
1801    async fn summarize_falls_back_to_heuristic_on_llm_error_by_default() {
1802        // Compatibility callers that do not provide exact model limits retain
1803        // the historical heuristic recovery behavior. Production compression
1804        // always supplies an explicit budget and surfaces stage failures.
1805        let summarizer =
1806            LlmSummarizer::new(Arc::new(FailingProvider), "model".to_string(), None, None);
1807        let out = summarizer.summarize(&summary_messages()).await;
1808        assert!(
1809            out.is_ok(),
1810            "default heuristic fallback should recover from a transient LLM error, got {out:?}"
1811        );
1812    }
1813
1814    #[tokio::test]
1815    async fn summarize_surfaces_llm_error_when_heuristic_fallback_disabled() {
1816        // Compatibility callers can still opt out of heuristic recovery.
1817        let summarizer =
1818            LlmSummarizer::new(Arc::new(FailingProvider), "model".to_string(), None, None)
1819                .with_heuristic_fallback_on_error(false);
1820        let out = summarizer.summarize(&summary_messages()).await;
1821        assert!(
1822            out.is_err(),
1823            "with the heuristic fallback disabled, a transient LLM error must surface"
1824        );
1825    }
1826
1827    #[derive(Default)]
1828    struct BoundedRequestCaptureProvider {
1829        requests: Mutex<Vec<(Vec<Message>, u32)>>,
1830    }
1831
1832    #[async_trait]
1833    impl LLMProvider for BoundedRequestCaptureProvider {
1834        async fn chat_stream(
1835            &self,
1836            messages: &[Message],
1837            _tools: &[bamboo_domain::ToolSchema],
1838            max_output_tokens: Option<u32>,
1839            _model: &str,
1840        ) -> Result<LLMStream, LLMError> {
1841            self.requests
1842                .lock()
1843                .expect("bounded request capture lock")
1844                .push((messages.to_vec(), max_output_tokens.unwrap_or_default()));
1845            Ok(Box::pin(stream::iter(vec![
1846                Ok::<LLMChunk, LLMError>(LLMChunk::Token(
1847                    "Detailed bounded summary with requirements, decisions, files, tests, and next steps. "
1848                        .repeat(16),
1849                )),
1850                Ok::<LLMChunk, LLMError>(LLMChunk::Done),
1851            ])))
1852        }
1853    }
1854
1855    #[derive(Default)]
1856    struct MultipartSharedContextProvider {
1857        requests: Mutex<Vec<(Vec<Message>, u32)>>,
1858    }
1859
1860    impl MultipartSharedContextProvider {
1861        fn echoed_shared_context(rendered: &str) -> String {
1862            let mut retained = vec!["SHARED_CAPSULE_763"];
1863            for sentinel in [
1864                "PREVIOUS_SENTINEL_763",
1865                "RUNTIME_CONTEXT_SENTINEL_763",
1866                "CUSTOM_INSTRUCTION_SENTINEL_763",
1867                "PRECOMPACT_SENTINEL_763",
1868            ] {
1869                if rendered.contains(sentinel) {
1870                    retained.push(sentinel);
1871                }
1872            }
1873            retained.join(" ")
1874        }
1875    }
1876
1877    #[async_trait]
1878    impl LLMProvider for MultipartSharedContextProvider {
1879        async fn chat_stream(
1880            &self,
1881            messages: &[Message],
1882            _tools: &[bamboo_domain::ToolSchema],
1883            max_output_tokens: Option<u32>,
1884            _model: &str,
1885        ) -> Result<LLMStream, LLMError> {
1886            self.requests
1887                .lock()
1888                .expect("multipart request capture lock")
1889                .push((messages.to_vec(), max_output_tokens.unwrap_or_default()));
1890            let rendered = messages
1891                .iter()
1892                .map(|message| message.content.as_str())
1893                .collect::<Vec<_>>()
1894                .join("\n");
1895            let response = if rendered.contains("shared-context capsule stage") {
1896                "SHARED_CAPSULE_763 PREVIOUS_SENTINEL_763 RUNTIME_CONTEXT_SENTINEL_763 \
1897                 CUSTOM_INSTRUCTION_SENTINEL_763 PRECOMPACT_SENTINEL_763"
1898                    .to_string()
1899            } else if rendered.contains("multipart final stage") {
1900                format!(
1901                    "FINAL_MULTIPART_SECTION_763 {}",
1902                    Self::echoed_shared_context(&rendered)
1903                )
1904            } else if rendered.contains("Preserve this shared context capsule")
1905                || rendered.contains("SHARED_CAPSULE_763")
1906                || rendered.contains("PREVIOUS_SENTINEL_763")
1907                || rendered.contains("RUNTIME_CONTEXT_SENTINEL_763")
1908                || rendered.contains("CUSTOM_INSTRUCTION_SENTINEL_763")
1909                || rendered.contains("PRECOMPACT_SENTINEL_763")
1910            {
1911                Self::echoed_shared_context(&rendered)
1912            } else if rendered.contains("intermediate reduce stage") {
1913                "INTERMEDIATE_PART_763 with detailed chronological evidence. ".repeat(12)
1914            } else {
1915                "MAP_PART_763 with detailed chronological evidence. ".repeat(12)
1916            };
1917            Ok(Box::pin(stream::iter(vec![
1918                Ok::<LLMChunk, LLMError>(LLMChunk::Token(response)),
1919                Ok::<LLMChunk, LLMError>(LLMChunk::Done),
1920            ])))
1921        }
1922    }
1923
1924    fn bounded_budget(
1925        context_window_tokens: u32,
1926        max_output_tokens: u32,
1927        safety_margin_tokens: u32,
1928        target_summary_tokens: u32,
1929    ) -> SummaryRequestBudget {
1930        let token_budget = TokenBudget::with_safety_margin(
1931            context_window_tokens,
1932            max_output_tokens,
1933            bamboo_compression::BudgetStrategy::default(),
1934            safety_margin_tokens,
1935        );
1936        SummaryRequestBudget::from_token_budget(&token_budget, 80, target_summary_tokens, 0.20)
1937    }
1938
1939    #[tokio::test]
1940    async fn bounded_small_source_still_uses_map_then_reduce() {
1941        let provider = Arc::new(BoundedRequestCaptureProvider::default());
1942        let budget = bounded_budget(10_000, 2_000, 100, 400);
1943        let summarizer =
1944            LlmSummarizer::new(provider.clone(), "summary-model".to_string(), None, None)
1945                .with_request_budget(budget.clone())
1946                .with_heuristic_fallback_on_error(false);
1947
1948        let report = summarizer
1949            .summarize_with_report(&summary_messages())
1950            .await
1951            .expect("bounded chunked summary");
1952        assert_eq!(report.map_calls, 1);
1953        assert_eq!(report.reduce_calls, 1);
1954
1955        let requests = provider.requests.lock().expect("capture lock");
1956        assert_eq!(requests.len(), 2);
1957        let map_request = requests
1958            .iter()
1959            .find(|(messages, _)| {
1960                messages
1961                    .iter()
1962                    .any(|message| message.content.contains("map stage"))
1963            })
1964            .expect("small input must still use a map request");
1965        let reduce_request = requests
1966            .iter()
1967            .find(|(messages, _)| {
1968                messages
1969                    .iter()
1970                    .any(|message| message.content.contains("final reduce stage"))
1971            })
1972            .expect("small input must still use a final reduce request");
1973        assert_eq!(reduce_request.1, 400);
1974        let counter = TiktokenTokenCounter::default();
1975        for (request, output) in [map_request, reduce_request] {
1976            assert!(
1977                counter
1978                    .count_messages(request)
1979                    .saturating_add(*output)
1980                    .saturating_add(budget.safety_margin_tokens)
1981                    <= budget.safe_request_tokens()
1982            );
1983        }
1984    }
1985
1986    #[tokio::test]
1987    async fn retained_shared_context_is_mapped_and_reduced_without_a_second_ratio() {
1988        let provider = Arc::new(BoundedRequestCaptureProvider::default());
1989        let budget = bounded_budget(10_000, 2_000, 100, 1_000);
1990        let summarizer =
1991            LlmSummarizer::new(provider.clone(), "summary-model".to_string(), None, None)
1992                .with_request_budget(budget.clone())
1993                .with_heuristic_fallback_on_error(false);
1994        let shared_context =
1995            "durable prior summary requirement decision exact path test evidence ".repeat(24);
1996        let shared_tokens = TiktokenTokenCounter::default().count_text(&shared_context);
1997
1998        let (_capsule, map_calls, reduce_calls) = summarizer
1999            .build_bounded_shared_context_capsule(&shared_context, shared_tokens, &budget)
2000            .await
2001            .expect("retained shared context should use bounded map/reduce");
2002
2003        assert_eq!(map_calls, 1);
2004        assert_eq!(reduce_calls, 1);
2005        let requests = provider.requests.lock().expect("capture lock");
2006        assert_eq!(requests.len(), 2);
2007        assert!(
2008            requests
2009                .iter()
2010                .all(|(_, output)| *output >= shared_tokens.saturating_sub(1)),
2011            "the child map/reduce must preserve its allocated prior-summary size instead of applying 20% again"
2012        );
2013    }
2014
2015    #[tokio::test]
2016    async fn bounded_pipeline_chunks_large_source_and_never_compounds_twenty_percent_target() {
2017        let provider = Arc::new(BoundedRequestCaptureProvider::default());
2018        let budget = bounded_budget(3_000, 800, 100, 600);
2019        let summarizer = LlmSummarizer::new(
2020            provider.clone(),
2021            "small-summary-model".to_string(),
2022            Some("Previous durable summary evidence. ".repeat(12)),
2023            None,
2024        )
2025        .with_context_blocks(vec![ContextBlock::new(
2026            ContextBlockType::TaskSnapshot,
2027            ContextBlockPriority::High,
2028            ContextBlockStability::RoundDynamic,
2029            "Current Task List",
2030            "Active compression task with exact acceptance evidence. ".repeat(8),
2031        )])
2032        .with_custom_instructions(Some(
2033            "Keep exact paths, failures, and remaining work.".to_string(),
2034        ))
2035        .with_request_budget(budget.clone())
2036        .with_heuristic_fallback_on_error(false);
2037        let messages = (0..80)
2038            .map(|index| {
2039                if index % 2 == 0 {
2040                    Message::user(format!(
2041                        "user-{index} {}",
2042                        "requirement decision path error evidence ".repeat(24)
2043                    ))
2044                } else {
2045                    Message::assistant(
2046                        format!(
2047                            "assistant-{index} {}",
2048                            "implementation command output test result next step ".repeat(24)
2049                        ),
2050                        None,
2051                    )
2052                }
2053            })
2054            .collect::<Vec<_>>();
2055
2056        let report = summarizer
2057            .summarize_with_report(&messages)
2058            .await
2059            .expect("chunked summary");
2060        assert!(report.map_calls > 1, "large source must use multiple maps");
2061        assert!(
2062            report.reduce_calls >= 1,
2063            "bounded partials should receive a final reduce when it fits"
2064        );
2065
2066        let requests = provider.requests.lock().expect("capture lock");
2067        let counter = TiktokenTokenCounter::default();
2068        for (request, output_tokens) in requests.iter() {
2069            assert!(
2070                counter
2071                    .count_messages(request)
2072                    .saturating_add(*output_tokens)
2073                    .saturating_add(budget.safety_margin_tokens)
2074                    <= budget.safe_request_tokens(),
2075                "every map/reduce request must satisfy the 80% invariant"
2076            );
2077            assert!(*output_tokens <= budget.max_output_tokens);
2078        }
2079        assert_eq!(
2080            requests.last().map(|(_, output)| *output),
2081            Some(600),
2082            "the final reduce keeps the global source-derived target instead of taking 20% of map summaries"
2083        );
2084        assert!(
2085            requests.iter().any(|(request, _)| request
2086                .iter()
2087                .any(|message| { message.content.contains("intermediate reduce stage") })),
2088            "large ordered partials should be recursively reduced before the final merge"
2089        );
2090        let final_request = requests.last().expect("final reduce request");
2091        let final_rendered = final_request
2092            .0
2093            .iter()
2094            .map(|message| message.content.as_str())
2095            .collect::<Vec<_>>()
2096            .join("\n");
2097        assert!(final_rendered.contains("Previous durable summary evidence"));
2098        assert!(final_rendered.contains("Current Task List"));
2099        assert!(final_rendered.contains("Keep exact paths"));
2100    }
2101
2102    #[tokio::test]
2103    async fn multipart_terminal_path_preserves_shared_context_and_instructions_boundedly() {
2104        let provider = Arc::new(MultipartSharedContextProvider::default());
2105        let budget = bounded_budget(3_000, 240, 100, 700);
2106        let summarizer = LlmSummarizer::new(
2107            provider.clone(),
2108            "small-summary-model".to_string(),
2109            Some(format!(
2110                "PREVIOUS_SENTINEL_763 durable prior state. {}",
2111                "large previous summary evidence ".repeat(2_000)
2112            )),
2113            None,
2114        )
2115        .with_context_blocks(vec![ContextBlock::new(
2116            ContextBlockType::TaskSnapshot,
2117            ContextBlockPriority::High,
2118            ContextBlockStability::RoundDynamic,
2119            "Current Runtime State",
2120            "RUNTIME_CONTEXT_SENTINEL_763 active objective and exact evidence.",
2121        )])
2122        .with_custom_instructions(Some(
2123            "CUSTOM_INSTRUCTION_SENTINEL_763 keep exact paths and failures.\n\n\
2124             ## PreCompact Hook Instructions\n\n\
2125             PRECOMPACT_SENTINEL_763 preserve the hook-injected next step."
2126                .to_string(),
2127        ))
2128        .with_request_budget(budget.clone())
2129        .with_heuristic_fallback_on_error(false);
2130        let messages = (0..100)
2131            .map(|index| {
2132                if index % 2 == 0 {
2133                    Message::user(format!(
2134                        "user-{index} {}",
2135                        "requirement decision path error evidence ".repeat(24)
2136                    ))
2137                } else {
2138                    Message::assistant(
2139                        format!(
2140                            "assistant-{index} {}",
2141                            "implementation command output test result next step ".repeat(24)
2142                        ),
2143                        None,
2144                    )
2145                }
2146            })
2147            .collect::<Vec<_>>();
2148
2149        let report = summarizer
2150            .summarize_with_report(&messages)
2151            .await
2152            .expect("multipart finalization should preserve shared context");
2153        assert!(report.map_calls > 1);
2154        assert!(
2155            report.reduce_calls >= 2,
2156            "the shared capsule and multipart final sections must be observable as reduce calls"
2157        );
2158        assert!(report.content.contains("PREVIOUS_SENTINEL_763"));
2159        assert!(report.content.contains("RUNTIME_CONTEXT_SENTINEL_763"));
2160        assert!(report.content.contains("CUSTOM_INSTRUCTION_SENTINEL_763"));
2161        assert!(report.content.contains("PRECOMPACT_SENTINEL_763"));
2162
2163        let requests = provider.requests.lock().expect("capture lock");
2164        let counter = TiktokenTokenCounter::default();
2165        for (request, output_tokens) in requests.iter() {
2166            assert!(
2167                counter
2168                    .count_messages(request)
2169                    .saturating_add(*output_tokens)
2170                    .saturating_add(budget.safety_margin_tokens)
2171                    <= budget.safe_request_tokens(),
2172                "every shared-capsule and multipart request must remain under the safe ceiling"
2173            );
2174        }
2175
2176        assert!(
2177            requests.iter().all(|(request, _)| !request
2178                .iter()
2179                .any(|message| message.content.contains("shared-context capsule stage"))),
2180            "large shared context must use the bounded hierarchical capsule path"
2181        );
2182        let capsule_inputs = requests
2183            .iter()
2184            .filter(|(request, _)| {
2185                request
2186                    .iter()
2187                    .any(|message| message.content.contains("You are the map stage"))
2188                    && request.iter().any(|message| {
2189                        message
2190                            .content
2191                            .contains("Preserve this shared context capsule")
2192                            || message.content.contains("PREVIOUS_SENTINEL_763")
2193                            || message.content.contains("RUNTIME_CONTEXT_SENTINEL_763")
2194                            || message.content.contains("CUSTOM_INSTRUCTION_SENTINEL_763")
2195                            || message.content.contains("PRECOMPACT_SENTINEL_763")
2196                    })
2197            })
2198            .flat_map(|(request, _)| request.iter())
2199            .map(|message| message.content.as_str())
2200            .collect::<Vec<_>>()
2201            .join("\n");
2202        assert!(capsule_inputs.contains("PREVIOUS_SENTINEL_763"));
2203        assert!(capsule_inputs.contains("RUNTIME_CONTEXT_SENTINEL_763"));
2204        assert!(capsule_inputs.contains("CUSTOM_INSTRUCTION_SENTINEL_763"));
2205        assert!(capsule_inputs.contains("PRECOMPACT_SENTINEL_763"));
2206
2207        let multipart_requests = requests
2208            .iter()
2209            .filter(|(request, _)| {
2210                request.iter().any(|message| {
2211                    message
2212                        .content
2213                        .contains("instruction-aware multipart final stage")
2214                })
2215            })
2216            .collect::<Vec<_>>();
2217        assert!(
2218            !multipart_requests.is_empty(),
2219            "global target above model output capacity must use multipart final sections"
2220        );
2221        for (request, _) in multipart_requests {
2222            let rendered = request
2223                .iter()
2224                .map(|message| message.content.as_str())
2225                .collect::<Vec<_>>()
2226                .join("\n");
2227            assert!(rendered.contains("SHARED_CAPSULE_763"));
2228            assert!(rendered.contains("PREVIOUS_SENTINEL_763"));
2229            assert!(rendered.contains("RUNTIME_CONTEXT_SENTINEL_763"));
2230            assert!(rendered.contains("CUSTOM_INSTRUCTION_SENTINEL_763"));
2231            assert!(rendered.contains("PRECOMPACT_SENTINEL_763"));
2232        }
2233    }
2234
2235    #[tokio::test]
2236    async fn multipart_terminal_path_without_shared_context_still_reduces_every_section() {
2237        let provider = Arc::new(BoundedRequestCaptureProvider::default());
2238        let budget = bounded_budget(3_000, 240, 100, 700);
2239        let summarizer = LlmSummarizer::new(
2240            provider.clone(),
2241            "small-summary-model".to_string(),
2242            None,
2243            None,
2244        )
2245        .with_request_budget(budget.clone())
2246        .with_heuristic_fallback_on_error(false);
2247        let messages = (0..100)
2248            .map(|index| {
2249                if index % 2 == 0 {
2250                    Message::user(format!(
2251                        "user-{index} {}",
2252                        "requirement decision path error evidence ".repeat(24)
2253                    ))
2254                } else {
2255                    Message::assistant(
2256                        format!(
2257                            "assistant-{index} {}",
2258                            "implementation command output test result next step ".repeat(24)
2259                        ),
2260                        None,
2261                    )
2262                }
2263            })
2264            .collect::<Vec<_>>();
2265
2266        let report = summarizer
2267            .summarize_with_report(&messages)
2268            .await
2269            .expect("multipart finalization should reduce every section");
2270        assert!(report.map_calls > 1);
2271        assert!(
2272            report.reduce_calls >= 1,
2273            "the terminal path must never persist unreduced map/intermediate partials"
2274        );
2275
2276        let requests = provider.requests.lock().expect("capture lock");
2277        assert!(requests.iter().any(|(request, _)| request
2278            .iter()
2279            .any(|message| message.content.contains("multipart final stage"))));
2280        let counter = TiktokenTokenCounter::default();
2281        for (request, output_tokens) in requests.iter() {
2282            assert!(
2283                counter
2284                    .count_messages(request)
2285                    .saturating_add(*output_tokens)
2286                    .saturating_add(budget.safety_margin_tokens)
2287                    <= budget.safe_request_tokens(),
2288                "every terminal reduce request must remain under the safe ceiling"
2289            );
2290        }
2291    }
2292
2293    #[test]
2294    fn oversized_terminal_partial_is_split_into_bounded_reduce_sections() {
2295        let budget = bounded_budget(3_000, 800, 100, 1_600);
2296        let summarizer = LlmSummarizer::new(
2297            Arc::new(DummyProvider),
2298            "small-summary-model".to_string(),
2299            None,
2300            None,
2301        )
2302        .with_request_budget(budget.clone());
2303        let counter = TiktokenTokenCounter::default();
2304        let capsule =
2305            counter.truncate_to_token_prefix(&"shared capsule evidence ".repeat(1_000), 575);
2306        let part = SummaryPart {
2307            content: counter
2308                .truncate_to_token_prefix(&"ordered partial evidence ".repeat(1_000), 800),
2309            represented_source_tokens: 4_000,
2310            first_message_id: "first".to_string(),
2311            last_message_id: "last".to_string(),
2312        };
2313        let requested_output = summarizer.target_for_source(part.represented_source_tokens);
2314        assert!(!summarizer.request_fits(
2315            &summarizer.build_multipart_finalize_messages(
2316                std::slice::from_ref(&part),
2317                requested_output,
2318                &capsule,
2319                true,
2320            ),
2321            requested_output,
2322            &budget,
2323        ));
2324
2325        let groups = summarizer
2326            .pack_multipart_final_groups(std::slice::from_ref(&part), &capsule, true, &budget)
2327            .expect("every finite terminal partial should split into bounded reduce sections");
2328        assert!(groups.len() > 1);
2329        assert_eq!(
2330            groups
2331                .iter()
2332                .flatten()
2333                .map(|part| part.represented_source_tokens)
2334                .sum::<u32>(),
2335            part.represented_source_tokens
2336        );
2337        for group in groups {
2338            let represented = group.iter().fold(0u32, |total, part| {
2339                total.saturating_add(part.represented_source_tokens)
2340            });
2341            let requested_output = summarizer.target_for_source(represented);
2342            assert!(summarizer.request_fits(
2343                &summarizer.build_multipart_finalize_messages(
2344                    &group,
2345                    requested_output,
2346                    &capsule,
2347                    true,
2348                ),
2349                requested_output,
2350                &budget,
2351            ));
2352        }
2353    }
2354
2355    #[test]
2356    fn one_hundred_thousand_raw_tokens_receive_twenty_thousand_token_target() {
2357        let summarizer = LlmSummarizer::new(
2358            Arc::new(DummyProvider),
2359            "summary-model".to_string(),
2360            None,
2361            None,
2362        )
2363        .with_request_budget(bounded_budget(128_000, 32_000, 1_000, 20_000));
2364
2365        assert_eq!(summarizer.target_for_source(100_000), 20_000);
2366    }
2367
2368    #[tokio::test]
2369    async fn bounded_stage_errors_never_fall_back_to_heuristic_by_default() {
2370        let summarizer = LlmSummarizer::new(
2371            Arc::new(FailingProvider),
2372            "summary-model".to_string(),
2373            None,
2374            None,
2375        )
2376        .with_request_budget(bounded_budget(10_000, 2_000, 100, 400));
2377
2378        let error = summarizer
2379            .summarize_with_report(&summary_messages())
2380            .await
2381            .expect_err("a failed bounded stage must surface atomically");
2382        assert!(error.to_string().contains("http 500 transient"));
2383    }
2384
2385    #[tokio::test]
2386    async fn hundreds_of_individually_small_messages_never_use_one_unbounded_request() {
2387        let provider = Arc::new(BoundedRequestCaptureProvider::default());
2388        let messages = (0..400)
2389            .map(|index| {
2390                if index % 2 == 0 {
2391                    Message::user(format!(
2392                        "small-user-{index} {}",
2393                        "requirement evidence detail ".repeat(4)
2394                    ))
2395                } else {
2396                    Message::assistant(
2397                        format!(
2398                            "small-assistant-{index} {}",
2399                            "result test next-step ".repeat(4)
2400                        ),
2401                        None,
2402                    )
2403                }
2404            })
2405            .collect::<Vec<_>>();
2406        let counter = TiktokenTokenCounter::default();
2407        let represented = counter.count_messages(&messages);
2408        let target = ((represented as f64) * 0.20).ceil() as u32;
2409        let budget = bounded_budget(3_000, 800, 100, target);
2410        let summarizer = LlmSummarizer::new(
2411            provider.clone(),
2412            "small-summary-model".to_string(),
2413            None,
2414            None,
2415        )
2416        .with_request_budget(budget.clone())
2417        .with_heuristic_fallback_on_error(false);
2418
2419        let report = summarizer
2420            .summarize_with_report(&messages)
2421            .await
2422            .expect("hundreds of small messages should chunk");
2423        assert!(report.map_calls > 1);
2424        assert_eq!(report.target_summary_tokens, target);
2425        let requests = provider.requests.lock().expect("capture lock");
2426        for (request, output) in requests.iter() {
2427            assert!(
2428                counter
2429                    .count_messages(request)
2430                    .saturating_add(*output)
2431                    .saturating_add(budget.safety_margin_tokens)
2432                    <= budget.safe_request_tokens()
2433            );
2434        }
2435        let raw_map_prompts = requests
2436            .iter()
2437            .filter(|(request, _)| {
2438                request
2439                    .iter()
2440                    .any(|message| message.content.contains("map stage"))
2441            })
2442            .flat_map(|(request, _)| request.iter())
2443            .map(|message| message.content.as_str())
2444            .collect::<Vec<_>>()
2445            .join("\n");
2446        assert!(raw_map_prompts.contains("small-user-0"));
2447        assert!(raw_map_prompts.contains("small-assistant-399"));
2448    }
2449
2450    #[tokio::test]
2451    async fn map_chunk_boundaries_do_not_split_generic_multi_tool_chains() {
2452        let provider = Arc::new(BoundedRequestCaptureProvider::default());
2453        let budget = bounded_budget(2_000, 300, 100, 300);
2454        let summarizer = LlmSummarizer::new(
2455            provider.clone(),
2456            "tiny-summary-model".to_string(),
2457            None,
2458            None,
2459        )
2460        .with_request_budget(budget)
2461        .with_heuristic_fallback_on_error(false);
2462        let mut messages = (0..8)
2463            .map(|index| {
2464                Message::user(format!(
2465                    "prefix-{index} {}",
2466                    "filler requirement evidence ".repeat(14)
2467                ))
2468            })
2469            .collect::<Vec<_>>();
2470        let mut chain = Message::assistant("CHAIN_ASSISTANT_763", None);
2471        chain.tool_calls = Some(vec![
2472            ToolCall {
2473                id: "chain-call-a-763".to_string(),
2474                tool_type: "function".to_string(),
2475                function: FunctionCall {
2476                    name: "search".to_string(),
2477                    arguments: r#"{"query":"a"}"#.to_string(),
2478                },
2479            },
2480            ToolCall {
2481                id: "chain-call-b-763".to_string(),
2482                tool_type: "function".to_string(),
2483                function: FunctionCall {
2484                    name: "read".to_string(),
2485                    arguments: r#"{"path":"b"}"#.to_string(),
2486                },
2487            },
2488        ]);
2489        messages.push(chain);
2490        messages.push(Message::tool_result(
2491            "chain-call-a-763",
2492            "CHAIN_RESULT_A_763",
2493        ));
2494        messages.push(Message::tool_result(
2495            "chain-call-b-763",
2496            "CHAIN_RESULT_B_763",
2497        ));
2498        messages.extend((0..8).map(|index| {
2499            Message::assistant(
2500                format!(
2501                    "suffix-{index} {}",
2502                    "implementation result evidence ".repeat(14)
2503                ),
2504                None,
2505            )
2506        }));
2507
2508        let report = summarizer
2509            .summarize_with_report(&messages)
2510            .await
2511            .expect("tool-chain source should chunk");
2512        assert!(report.map_calls > 1);
2513        let requests = provider.requests.lock().expect("capture lock");
2514        let chain_map = requests
2515            .iter()
2516            .flat_map(|(request, _)| request.iter())
2517            .find(|message| message.content.contains("CHAIN_RESULT_A_763"))
2518            .expect("map request containing tool chain");
2519        assert!(chain_map.content.contains("CHAIN_ASSISTANT_763"));
2520        assert!(chain_map.content.contains("CHAIN_RESULT_B_763"));
2521    }
2522
2523    #[tokio::test]
2524    async fn oversized_single_message_is_split_without_dropping_its_tail() {
2525        let provider = Arc::new(BoundedRequestCaptureProvider::default());
2526        let budget = bounded_budget(2_000, 300, 100, 300);
2527        let summarizer = LlmSummarizer::new(
2528            provider.clone(),
2529            "tiny-summary-model".to_string(),
2530            None,
2531            None,
2532        )
2533        .with_request_budget(budget)
2534        .with_heuristic_fallback_on_error(false);
2535        let messages = vec![Message::user(format!(
2536            "{} TAIL_SENTINEL_763",
2537            "very large source message with concrete content ".repeat(2_000)
2538        ))];
2539
2540        let report = summarizer
2541            .summarize_with_report(&messages)
2542            .await
2543            .expect("oversized source should split");
2544        assert!(report.map_calls > 1);
2545        let requests = provider.requests.lock().expect("capture lock");
2546        assert!(
2547            requests.iter().any(|(request, _)| request
2548                .iter()
2549                .any(|message| message.content.contains("TAIL_SENTINEL_763"))),
2550            "the deterministic continuation chunks must include the original tail"
2551        );
2552    }
2553
2554    struct PartialWithoutDoneProvider;
2555
2556    #[async_trait]
2557    impl LLMProvider for PartialWithoutDoneProvider {
2558        async fn chat_stream(
2559            &self,
2560            _messages: &[Message],
2561            _tools: &[bamboo_domain::ToolSchema],
2562            _max_output_tokens: Option<u32>,
2563            _model: &str,
2564        ) -> Result<LLMStream, LLMError> {
2565            Ok(Box::pin(stream::iter(vec![Ok::<LLMChunk, LLMError>(
2566                LLMChunk::Token("partial but incomplete summary".to_string()),
2567            )])))
2568        }
2569    }
2570
2571    #[tokio::test]
2572    async fn partial_stream_without_done_is_never_accepted_as_summary() {
2573        let summarizer = LlmSummarizer::new(
2574            Arc::new(PartialWithoutDoneProvider),
2575            "model".to_string(),
2576            None,
2577            None,
2578        )
2579        .with_request_budget(bounded_budget(10_000, 2_000, 100, 400))
2580        .with_heuristic_fallback_on_error(false);
2581        let error = summarizer
2582            .summarize_with_report(&summary_messages())
2583            .await
2584            .expect_err("partial stream must fail");
2585        assert!(error.to_string().contains("without terminal completion"));
2586    }
2587}