1use 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#[derive(Debug, Clone, Default)]
36pub enum SummaryMode {
37 #[default]
39 FullRewrite,
40 IncrementalMerge,
42}
43
44#[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
127pub struct LlmSummarizer {
136 llm: Arc<dyn LLMProvider>,
137 model: String,
138 existing_summary: Option<String>,
140 context_blocks: Vec<ContextBlock>,
142 custom_instructions: Option<String>,
144 summary_mode: SummaryMode,
146 heuristic_fallback_on_error: bool,
152 request_budget: Option<SummaryRequestBudget>,
154 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 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 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 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 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 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 async fn collect_stream_response(
1343 &self,
1344 messages: &[Message],
1345 max_output_tokens: u32,
1346 ) -> Result<String, bamboo_compression::types::BudgetError> {
1347 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(_) => {} 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 #[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 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 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 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}