kmp_application/queries/
bundle_truncator.rs

1use kmp_domain::{KmpMode, ResolutionTier, TierBudget, TokenEstimator};
2
3use super::bundle_section_renderer::TaggedSection;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct TruncationMetadata {
7    pub budget_requested: u32,
8    pub budget_used: u32,
9    pub total_before_truncation: u32,
10    pub sections_kept: u32,
11    pub sections_dropped: u32,
12    pub token_estimator: String,
13}
14
15/// Tier-aware truncation: L0 guaranteed, L1 prioritized, L2 sacrificed.
16///
17/// Unlike greedy sequential packing, this continues past L2 sections that
18/// don't fit — later L1 sections can still be included if their tier budget
19/// allows. Returns `(content, source_id)` pairs and truncation metadata.
20pub(crate) fn limit_sections_by_tier_budget(
21    sections: Vec<TaggedSection>,
22    token_budget: Option<u32>,
23    resolved_mode: KmpMode,
24    estimator: &dyn TokenEstimator,
25) -> (Vec<(String, String)>, Option<TruncationMetadata>) {
26    let Some(budget) = token_budget else {
27        let pairs = sections
28            .into_iter()
29            .map(|s| (s.content, s.source_id))
30            .collect();
31        return (pairs, None);
32    };
33
34    let tier_budget = TierBudget::from_total_with_mode(budget, resolved_mode);
35    let total_sections = sections.len() as u32;
36    let total_before: u32 = sections
37        .iter()
38        .map(|s| estimator.estimate_tokens(&s.content))
39        .sum();
40
41    let mut l0_used = 0u32;
42    let mut l1_used = 0u32;
43    let mut l2_used = 0u32;
44    let mut kept = Vec::new();
45    let mut total_used = 0u32;
46
47    for section in sections {
48        let tokens = estimator.estimate_tokens(&section.content);
49        let (tier_used, tier_cap) = match section.tier {
50            ResolutionTier::L0Summary => (&mut l0_used, tier_budget.l0),
51            ResolutionTier::L1CausalSpine => (&mut l1_used, tier_budget.l1),
52            ResolutionTier::L2EvidencePack => (&mut l2_used, tier_budget.l2),
53        };
54
55        // First section always included (L0 anchor). Otherwise check tier + total budget.
56        let fits_tier = kept.is_empty() || *tier_used + tokens <= tier_cap;
57        let fits_total = kept.is_empty() || total_used + tokens <= budget;
58
59        if fits_tier && fits_total {
60            *tier_used += tokens;
61            total_used += tokens;
62            kept.push((section.content, section.source_id));
63        }
64        // Don't break — a later section from a different tier may still fit.
65    }
66
67    let sections_kept = kept.len() as u32;
68    let truncation = TruncationMetadata {
69        budget_requested: budget,
70        budget_used: total_used,
71        total_before_truncation: total_before,
72        sections_kept,
73        sections_dropped: total_sections - sections_kept,
74        token_estimator: estimator.name().to_string(),
75    };
76
77    (kept, Some(truncation))
78}
kmp_application/queries/bundle_truncator.rs

kmp_application/queries/
bundle_truncator.rs
