bamboo_compression/
types.rs

1//! Core types for token budget management.
2//!
3//! TokenBudget, BudgetStrategy, TokenUsageBreakdown, and TokenBudgetUsage
4//! are re-exported from bamboo-domain-session.
5//! This file keeps PreparedContext and BudgetError in the facade for Phase 1.
6
7// Re-exported from domain crate
8pub use bamboo_domain::budget_types::{
9    BudgetStrategy, TokenBudget, TokenBudgetUsage, TokenUsageBreakdown,
10};
11
12use thiserror::Error;
13
14/// Result of context preparation with budget enforcement.
15#[derive(Debug, Clone)]
16pub struct PreparedContext {
17    /// Messages prepared for LLM (may be truncated)
18    pub messages: Vec<bamboo_domain::Message>,
19    /// Token usage breakdown
20    pub token_usage: TokenUsageBreakdown,
21    /// Whether truncation occurred
22    pub truncation_occurred: bool,
23    /// Number of message segments removed
24    pub segments_removed: usize,
25    /// Message IDs newly archived by this preparation pass.
26    pub compressed_message_ids: Vec<String>,
27    /// Number of long tool outputs replaced with prompt-side cached summaries.
28    pub prompt_cached_tool_outputs: usize,
29    /// Tokens saved by prompt-side tool output compaction in this preparation pass.
30    pub prompt_cached_tool_tokens_saved: u32,
31}
32
33/// Errors that can occur during budget management.
34#[derive(Debug, Error)]
35pub enum BudgetError {
36    #[error("System prompt ({system_tokens} tokens) exceeds available budget ({available_tokens} tokens)")]
37    SystemPromptTooLarge {
38        system_tokens: u32,
39        available_tokens: u32,
40    },
41
42    #[error("Single message ({message_tokens} tokens) exceeds available budget ({available_tokens} tokens). Consider splitting the message or attaching as a file.")]
43    SingleMessageTooLarge {
44        message_tokens: u32,
45        available_tokens: u32,
46    },
47
48    #[error("Failed to count tokens: {0}")]
49    TokenCountError(String),
50
51    #[error("Failed to segment messages: {0}")]
52    SegmentationError(String),
53}
54
55#[cfg(test)]
56mod tests {
57    use super::{BudgetStrategy, TokenBudget};
58
59    #[test]
60    fn compression_trigger_defaults_to_eighty_five_percent() {
61        let budget = TokenBudget::for_model(128_000);
62        assert_eq!(budget.compression_trigger_percent, 85);
63    }
64
65    #[test]
66    fn compression_target_defaults_to_forty_percent() {
67        let budget = TokenBudget::for_model(128_000);
68        assert_eq!(budget.compression_target_percent, 40);
69    }
70
71    #[test]
72    fn prompt_cache_defaults_match_current_compaction_policy() {
73        let budget = TokenBudget::for_model(128_000);
74        assert_eq!(budget.prompt_cache_min_tool_output_chars, 1_200);
75        assert_eq!(budget.prompt_cache_head_chars, 280);
76        assert_eq!(budget.prompt_cache_tail_chars, 180);
77        assert_eq!(budget.prompt_cache_recent_user_turns, 2);
78        assert_eq!(budget.prompt_cache_recent_tool_chains, 2);
79    }
80
81    #[test]
82    fn compression_trigger_context_tokens_respects_percent() {
83        let mut budget =
84            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
85        budget.working_reserve_tokens = 0; // use legacy percentage mode
86        budget.compression_trigger_percent = 50;
87        assert_eq!(budget.compression_trigger_context_tokens(), 500);
88    }
89
90    #[test]
91    fn compression_target_context_tokens_respects_percent() {
92        let mut budget =
93            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
94        budget.compression_target_percent = 50;
95        assert_eq!(budget.compression_target_context_tokens(), 500);
96    }
97
98    #[test]
99    fn compression_target_percent_is_clamped_to_supported_range() {
100        let mut budget =
101            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
102        budget.compression_target_percent = 10;
103        assert_eq!(budget.compression_target_context_tokens(), 200);
104    }
105
106    #[test]
107    fn compression_target_always_stays_below_trigger_limit() {
108        let mut budget =
109            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
110        budget.working_reserve_tokens = 0; // use legacy percentage mode
111        budget.compression_trigger_percent = 30;
112        budget.compression_target_percent = 50;
113        assert_eq!(budget.compression_target_context_tokens(), 299);
114    }
115
116    #[test]
117    fn trigger_percent_zero_means_disabled() {
118        let mut budget =
119            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
120        budget.working_reserve_tokens = 0; // use legacy percentage mode
121        budget.compression_trigger_percent = 0;
122        assert_eq!(
123            budget.compression_trigger_context_tokens(),
124            budget.max_context_tokens
125        );
126    }
127
128    #[test]
129    fn fixed_reserve_trigger_subtracts_reserve_from_context() {
130        let budget =
131            TokenBudget::with_safety_margin(200_000, 4_096, BudgetStrategy::default(), 1000);
132        // default working_reserve_tokens = 50_000
133        assert_eq!(budget.compression_trigger_context_tokens(), 150_000);
134    }
135
136    #[test]
137    fn fixed_reserve_trigger_for_small_context() {
138        let budget =
139            TokenBudget::with_safety_margin(100_000, 4_096, BudgetStrategy::default(), 1000);
140        // 100K >= 50K * 2, so fixed reserve: 100K - 50K = 50K
141        assert_eq!(budget.compression_trigger_context_tokens(), 50_000);
142    }
143
144    #[test]
145    fn fixed_reserve_fallback_for_tiny_context() {
146        let budget =
147            TokenBudget::with_safety_margin(60_000, 4_096, BudgetStrategy::default(), 1000);
148        // 60K < 50K * 2 = 100K, so fallback to 75%
149        assert_eq!(budget.compression_trigger_context_tokens(), 45_000);
150    }
151
152    #[test]
153    fn working_reserve_zero_uses_legacy_percentage() {
154        let mut budget =
155            TokenBudget::with_safety_margin(200_000, 4_096, BudgetStrategy::default(), 1000);
156        budget.working_reserve_tokens = 0;
157        budget.compression_trigger_percent = 85;
158        assert_eq!(budget.compression_trigger_context_tokens(), 170_000);
159    }
160}
bamboo_compression/types.rs

bamboo_compression/
types.rs
