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}
30
31/// Errors that can occur during budget management.
32#[derive(Debug, Error)]
33pub enum BudgetError {
34    #[error("System prompt ({system_tokens} tokens) exceeds available budget ({available_tokens} tokens)")]
35    SystemPromptTooLarge {
36        system_tokens: u32,
37        available_tokens: u32,
38    },
39
40    #[error("Single message ({message_tokens} tokens) exceeds available budget ({available_tokens} tokens). Consider splitting the message or attaching as a file.")]
41    SingleMessageTooLarge {
42        message_tokens: u32,
43        available_tokens: u32,
44    },
45
46    #[error("Failed to count tokens: {0}")]
47    TokenCountError(String),
48
49    #[error("Failed to segment messages: {0}")]
50    SegmentationError(String),
51}
52
53#[cfg(test)]
54mod tests {
55    use super::{BudgetStrategy, TokenBudget};
56
57    #[test]
58    fn compression_trigger_defaults_to_eighty_five_percent() {
59        let budget = TokenBudget::for_model(128_000);
60        assert_eq!(budget.compression_trigger_percent, 85);
61    }
62
63    #[test]
64    fn compression_target_defaults_to_forty_percent() {
65        let budget = TokenBudget::for_model(128_000);
66        assert_eq!(budget.compression_target_percent, 40);
67    }
68
69    #[test]
70    fn prompt_cache_defaults_match_current_compaction_policy() {
71        let budget = TokenBudget::for_model(128_000);
72        assert_eq!(budget.prompt_cache_min_tool_output_chars, 1_200);
73        assert_eq!(budget.prompt_cache_head_chars, 280);
74        assert_eq!(budget.prompt_cache_tail_chars, 180);
75        assert_eq!(budget.prompt_cache_recent_user_turns, 2);
76        assert_eq!(budget.prompt_cache_recent_tool_chains, 2);
77    }
78
79    #[test]
80    fn compression_trigger_context_tokens_respects_percent() {
81        let mut budget =
82            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
83        budget.working_reserve_tokens = 0; // use legacy percentage mode
84        budget.compression_trigger_percent = 50;
85        assert_eq!(budget.compression_trigger_context_tokens(), 500);
86    }
87
88    #[test]
89    fn compression_target_context_tokens_respects_percent() {
90        let mut budget =
91            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
92        budget.compression_target_percent = 50;
93        assert_eq!(budget.compression_target_context_tokens(), 500);
94    }
95
96    #[test]
97    fn compression_target_percent_is_clamped_to_supported_range() {
98        let mut budget =
99            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
100        budget.compression_target_percent = 10;
101        assert_eq!(budget.compression_target_context_tokens(), 200);
102    }
103
104    #[test]
105    fn compression_target_always_stays_below_trigger_limit() {
106        let mut budget =
107            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
108        budget.working_reserve_tokens = 0; // use legacy percentage mode
109        budget.compression_trigger_percent = 30;
110        budget.compression_target_percent = 50;
111        assert_eq!(budget.compression_target_context_tokens(), 299);
112    }
113
114    #[test]
115    fn trigger_percent_zero_means_disabled() {
116        let mut budget =
117            TokenBudget::with_safety_margin(1000, 200, BudgetStrategy::Window { size: 20 }, 100);
118        budget.working_reserve_tokens = 0; // use legacy percentage mode
119        budget.compression_trigger_percent = 0;
120        assert_eq!(
121            budget.compression_trigger_context_tokens(),
122            budget.max_context_tokens
123        );
124    }
125
126    #[test]
127    fn fixed_reserve_trigger_subtracts_reserve_from_context() {
128        let budget =
129            TokenBudget::with_safety_margin(200_000, 4_096, BudgetStrategy::default(), 1000);
130        // default working_reserve_tokens = 50_000
131        assert_eq!(budget.compression_trigger_context_tokens(), 150_000);
132    }
133
134    #[test]
135    fn fixed_reserve_trigger_for_small_context() {
136        let budget =
137            TokenBudget::with_safety_margin(100_000, 4_096, BudgetStrategy::default(), 1000);
138        // 100K >= 50K * 2, so fixed reserve: 100K - 50K = 50K
139        assert_eq!(budget.compression_trigger_context_tokens(), 50_000);
140    }
141
142    #[test]
143    fn fixed_reserve_fallback_for_tiny_context() {
144        let budget =
145            TokenBudget::with_safety_margin(60_000, 4_096, BudgetStrategy::default(), 1000);
146        // 60K < 50K * 2 = 100K, so fallback to 75%
147        assert_eq!(budget.compression_trigger_context_tokens(), 45_000);
148    }
149
150    #[test]
151    fn working_reserve_zero_uses_legacy_percentage() {
152        let mut budget =
153            TokenBudget::with_safety_margin(200_000, 4_096, BudgetStrategy::default(), 1000);
154        budget.working_reserve_tokens = 0;
155        budget.compression_trigger_percent = 85;
156        assert_eq!(budget.compression_trigger_context_tokens(), 170_000);
157    }
158}
bamboo_compression/types.rs

bamboo_compression/
types.rs
