lean_ctx/core/
adaptive_compression.rs

1//! Cache-aware compression decisions using provider billing economics (#1195).
2//!
3//! Compression of an already-sent prompt prefix can turn cheap cache reads into
4//! cache writes.  This module compares that penalty with the actual token
5//! reduction and chooses the candidate with the highest expected USD benefit.
6
7use crate::core::gain::model_pricing::{ModelPricing, PricingMatchKind};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum CompressionDepth {
11    Conservative,
12    Balanced,
13    Aggressive,
14}
15
16#[derive(Debug, Clone, Copy)]
17pub struct CompressionCandidate {
18    pub depth: CompressionDepth,
19    pub compressed_tokens: u64,
20    /// Probability that the already-sent prefix remains byte-stable, in 0..=1.
21    pub prefix_stability: f64,
22}
23
24#[derive(Debug, Clone, Copy)]
25pub struct CompressionContext<'a> {
26    pub original_tokens: u64,
27    /// Tokens in the cacheable prefix before compression.
28    pub prefix_tokens: u64,
29    pub current_turn: u64,
30    pub predicted_total_turns: u64,
31    /// Measured probability of a compressed read causing an expensive retry.
32    pub bounce_rate: f64,
33    pub model: Option<&'a str>,
34}
35
36#[derive(Debug, Clone, Copy)]
37pub struct CandidateEconomics {
38    pub depth: CompressionDepth,
39    pub gross_benefit_usd: f64,
40    pub cache_break_risk_usd: f64,
41    pub net_benefit_usd: f64,
42    pub cache_break_probability: f64,
43}
44
45#[derive(Debug, Clone)]
46pub struct CompressionDecision {
47    pub selected: Option<CandidateEconomics>,
48    pub pricing_match: PricingMatchKind,
49}
50
51impl CompressionDecision {
52    #[must_use]
53    pub fn should_compress(&self) -> bool {
54        self.selected.is_some()
55    }
56}
57
58/// Compare candidates using all provider input token classes affected here.
59///
60/// `gross_benefit` prices removed tokens once as `cache_write`, then as
61/// `cache_read` on every predicted reuse. `cache_break_risk` prices the
62/// remaining compressed prefix at the write-minus-read premium on each reuse,
63/// weighted by the measured probability that the prefix changes.
64#[must_use]
65pub fn decide(
66    context: CompressionContext<'_>,
67    candidates: &[CompressionCandidate],
68) -> CompressionDecision {
69    let quote = ModelPricing::load().quote(context.model);
70    let remaining_turns = context
71        .predicted_total_turns
72        .saturating_sub(context.current_turn);
73    let write_rate = quote.cost.cache_write_per_m;
74    let read_rate = quote.cost.cache_read_per_m;
75    let write_premium = (write_rate - read_rate).max(0.0);
76
77    let selected = candidates
78        .iter()
79        .filter(|candidate| candidate.compressed_tokens < context.original_tokens)
80        .map(|candidate| {
81            let saved = context
82                .original_tokens
83                .saturating_sub(candidate.compressed_tokens);
84            let stability = finite_probability(candidate.prefix_stability);
85            let bounce_rate = finite_probability(context.bounce_rate);
86            let break_probability = 1.0 - stability * (1.0 - bounce_rate);
87            let gross_rate = write_rate + remaining_turns as f64 * read_rate;
88            let gross_benefit_usd = usd(saved, gross_rate);
89            let compressed_prefix = context.prefix_tokens.saturating_sub(saved);
90            let cache_break_risk_usd = usd(
91                compressed_prefix,
92                remaining_turns as f64 * break_probability * write_premium,
93            );
94            CandidateEconomics {
95                depth: candidate.depth,
96                gross_benefit_usd,
97                cache_break_risk_usd,
98                net_benefit_usd: gross_benefit_usd - cache_break_risk_usd,
99                cache_break_probability: break_probability,
100            }
101        })
102        .filter(|economics| economics.net_benefit_usd > 0.0)
103        .max_by(|a, b| {
104            a.net_benefit_usd
105                .total_cmp(&b.net_benefit_usd)
106                .then_with(|| depth_rank(a.depth).cmp(&depth_rank(b.depth)))
107        });
108
109    CompressionDecision {
110        selected,
111        pricing_match: quote.match_kind,
112    }
113}
114
115/// Ratio of the previous prefix preserved at the start of the current prefix.
116/// Appending content is perfectly stable; a mutation near the front scores low.
117#[must_use]
118pub fn prefix_stability_score(previous: &[u8], current: &[u8]) -> f64 {
119    if previous.is_empty() {
120        return 1.0;
121    }
122    let common = previous
123        .iter()
124        .zip(current)
125        .take_while(|(a, b)| a == b)
126        .count();
127    common as f64 / previous.len() as f64
128}
129
130/// Deterministic session-length estimate used when no explicit forecast exists.
131/// Short sessions retain a small horizon; established long sessions reserve a
132/// further 50% of their current length for cache-reuse economics.
133#[must_use]
134pub fn predict_session_length(current_turn: u64) -> u64 {
135    match current_turn {
136        0..=3 => 4,
137        4..=11 => 12,
138        _ => current_turn.saturating_add(current_turn / 2),
139    }
140}
141
142fn finite_probability(value: f64) -> f64 {
143    if value.is_finite() {
144        value.clamp(0.0, 1.0)
145    } else {
146        0.0
147    }
148}
149
150fn usd(tokens: u64, rate_per_m: f64) -> f64 {
151    tokens as f64 / 1_000_000.0 * rate_per_m
152}
153
154fn depth_rank(depth: CompressionDepth) -> u8 {
155    match depth {
156        CompressionDepth::Conservative => 0,
157        CompressionDepth::Balanced => 1,
158        CompressionDepth::Aggressive => 2,
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    fn context(turn: u64, predicted: u64) -> CompressionContext<'static> {
167        CompressionContext {
168            original_tokens: 1_000_000,
169            prefix_tokens: 1_000_000,
170            current_turn: turn,
171            predicted_total_turns: predicted,
172            bounce_rate: 0.0,
173            model: Some("claude-opus-4.5"),
174        }
175    }
176
177    #[test]
178    fn stable_prefix_prices_write_and_future_reads_exactly() {
179        let decision = decide(
180            context(1, 2),
181            &[CompressionCandidate {
182                depth: CompressionDepth::Balanced,
183                compressed_tokens: 500_000,
184                prefix_stability: 1.0,
185            }],
186        );
187        let economics = decision.selected.unwrap();
188        assert!((economics.gross_benefit_usd - 3.375).abs() < 1e-12);
189        assert_eq!(economics.cache_break_risk_usd, 0.0);
190        assert!((economics.net_benefit_usd - 3.375).abs() < 1e-12);
191    }
192
193    #[test]
194    fn unstable_long_prefix_costs_real_write_minus_read_premium() {
195        let decision = decide(
196            context(1, 13),
197            &[CompressionCandidate {
198                depth: CompressionDepth::Balanced,
199                compressed_tokens: 500_000,
200                prefix_stability: 0.0,
201            }],
202        );
203        assert!(!decision.should_compress());
204        let risk = 0.5 * 12.0 * (6.25 - 0.50);
205        assert!((risk - 34.5_f64).abs() < 1e-12);
206    }
207
208    #[test]
209    fn short_session_selects_largest_real_saving() {
210        let candidates = [
211            CompressionCandidate {
212                depth: CompressionDepth::Conservative,
213                compressed_tokens: 800_000,
214                prefix_stability: 0.98,
215            },
216            CompressionCandidate {
217                depth: CompressionDepth::Aggressive,
218                compressed_tokens: 300_000,
219                prefix_stability: 0.95,
220            },
221        ];
222        let selected = decide(context(3, 4), &candidates).selected.unwrap();
223        assert_eq!(selected.depth, CompressionDepth::Aggressive);
224    }
225
226    #[test]
227    fn long_session_can_prefer_conservative_stable_candidate() {
228        let candidates = [
229            CompressionCandidate {
230                depth: CompressionDepth::Conservative,
231                compressed_tokens: 800_000,
232                prefix_stability: 0.999,
233            },
234            CompressionCandidate {
235                depth: CompressionDepth::Aggressive,
236                compressed_tokens: 300_000,
237                prefix_stability: 0.75,
238            },
239        ];
240        let selected = decide(context(20, 40), &candidates).selected.unwrap();
241        assert_eq!(selected.depth, CompressionDepth::Conservative);
242    }
243
244    #[test]
245    fn non_saving_candidate_never_compresses() {
246        let candidate = CompressionCandidate {
247            depth: CompressionDepth::Aggressive,
248            compressed_tokens: 1_000_000,
249            prefix_stability: 1.0,
250        };
251        assert!(!decide(context(1, 1), &[candidate]).should_compress());
252    }
253
254    #[test]
255    fn appended_prefix_is_fully_stable() {
256        assert_eq!(prefix_stability_score(b"stable", b"stable + new"), 1.0);
257    }
258
259    #[test]
260    fn front_mutation_has_zero_stability() {
261        assert_eq!(prefix_stability_score(b"abc", b"xbc"), 0.0);
262    }
263
264    #[test]
265    fn prediction_distinguishes_short_and_long_sessions() {
266        assert_eq!(predict_session_length(2), 4);
267        assert_eq!(predict_session_length(8), 12);
268        assert_eq!(predict_session_length(40), 60);
269    }
270
271    #[test]
272    fn nan_stability_fails_closed_as_cache_break() {
273        let candidate = CompressionCandidate {
274            depth: CompressionDepth::Balanced,
275            compressed_tokens: 500_000,
276            prefix_stability: f64::NAN,
277        };
278        assert!(!decide(context(1, 13), &[candidate]).should_compress());
279    }
280
281    #[test]
282    fn measured_bounces_raise_expected_cache_risk() {
283        let candidate = CompressionCandidate {
284            depth: CompressionDepth::Balanced,
285            compressed_tokens: 500_000,
286            prefix_stability: 1.0,
287        };
288        let clean = decide(context(1, 13), &[candidate]).selected.unwrap();
289        let mut retrying = context(1, 13);
290        retrying.bounce_rate = 0.5;
291        assert!(decide(retrying, &[candidate]).selected.is_none());
292        assert_eq!(clean.cache_break_probability, 0.0);
293    }
294}
lean_ctx/core/adaptive_compression.rs

lean_ctx/core/
adaptive_compression.rs
