lean_ctx/core/
budget.rs

1//! Turn-level context budget controller (#1306).
2//!
3//! Caps fresh tokens per tool response to prevent context window bloat.
4//! Research basis: "Context Length Alone Hurts LLM Performance" (EMNLP 2025)
5//! found 13.9–85% degradation with length even with perfect retrieval.
6
7use super::tokens::count_tokens;
8
9/// Budget enforcement result.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum BudgetAction {
12    /// Content fits within budget — pass through unchanged.
13    PassThrough,
14    /// Content exceeds budget — truncated with expand hint appended.
15    Truncated {
16        original_tokens: usize,
17        delivered_tokens: usize,
18    },
19}
20
21/// Apply turn-level token budget to a tool response body.
22///
23/// If the response exceeds `fresh_limit` tokens, truncates to fit and appends
24/// an expand hint so the agent can retrieve the remainder.
25///
26/// Returns `(possibly_truncated_text, action)`.
27pub fn apply_turn_budget(text: &str, fresh_limit: usize) -> (String, BudgetAction) {
28    if fresh_limit == 0 {
29        return (text.to_string(), BudgetAction::PassThrough);
30    }
31
32    let token_count = count_tokens(text);
33    if token_count <= fresh_limit {
34        return (text.to_string(), BudgetAction::PassThrough);
35    }
36
37    let truncated = truncate_to_token_budget(text, fresh_limit);
38    let delivered_tokens = count_tokens(&truncated);
39
40    let hint = format!(
41        "\n[… truncated at ~{delivered_tokens} of {token_count} tokens — \
42         use ctx_read with lines= parameter to see specific sections]"
43    );
44
45    (
46        format!("{truncated}{hint}"),
47        BudgetAction::Truncated {
48            original_tokens: token_count,
49            delivered_tokens,
50        },
51    )
52}
53
54/// Truncate text to approximately `limit` tokens by keeping complete lines.
55fn truncate_to_token_budget(text: &str, limit: usize) -> String {
56    let mut result = String::new();
57    let mut current_tokens = 0;
58
59    for line in text.lines() {
60        let line_tokens = count_tokens(line);
61        if current_tokens + line_tokens > limit && current_tokens > 0 {
62            break;
63        }
64        if !result.is_empty() {
65            result.push('\n');
66        }
67        result.push_str(line);
68        current_tokens += line_tokens;
69    }
70
71    result
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn passthrough_when_within_budget() {
80        let text = "small content";
81        let (result, action) = apply_turn_budget(text, 1000);
82        assert_eq!(result, text);
83        assert_eq!(action, BudgetAction::PassThrough);
84    }
85
86    #[test]
87    fn passthrough_when_budget_is_zero() {
88        let text = "any content at all";
89        let (result, action) = apply_turn_budget(text, 0);
90        assert_eq!(result, text);
91        assert_eq!(action, BudgetAction::PassThrough);
92    }
93
94    #[test]
95    fn truncates_large_content() {
96        let lines: Vec<String> = (0..200)
97            .map(|i| format!("fn function_{i}() {{ let x = {i}; }}"))
98            .collect();
99        let text = lines.join("\n");
100        let (result, action) = apply_turn_budget(&text, 100);
101
102        assert!(result.contains("[… truncated"));
103        assert!(result.contains("use ctx_read with lines="));
104        match action {
105            BudgetAction::Truncated {
106                original_tokens,
107                delivered_tokens,
108            } => {
109                assert!(
110                    delivered_tokens <= 120,
111                    "delivered {delivered_tokens} > ~120"
112                );
113                assert!(original_tokens > delivered_tokens);
114            }
115            BudgetAction::PassThrough => panic!("should have truncated"),
116        }
117    }
118
119    #[test]
120    fn truncation_preserves_complete_lines() {
121        let text = "line one\nline two\nline three\nline four\nline five";
122        let (result, _) = apply_turn_budget(text, 5);
123        let body = result.split("\n[… truncated").next().unwrap();
124        assert!(
125            !body.ends_with(char::is_whitespace),
126            "truncated body should end with a complete line"
127        );
128    }
129}
lean_ctx/core/budget.rs

lean_ctx/core/
budget.rs
