stakpak_server/context/
budget.rs

1use crate::context::project::{ContextFile, ContextPriority};
2
3const DEFAULT_HEAD_RATIO: f64 = 0.7;
4const DEFAULT_TAIL_RATIO: f64 = 0.2;
5const MIN_FILE_ALLOCATION_CHARS: usize = 64;
6
7#[derive(Debug, Clone)]
8pub struct ContextBudget {
9    pub system_prompt_max_chars: usize,
10    pub per_file_max_chars: usize,
11    pub total_context_max_chars: usize,
12    pub head_ratio: f64,
13    pub tail_ratio: f64,
14}
15
16impl Default for ContextBudget {
17    fn default() -> Self {
18        Self {
19            system_prompt_max_chars: 32_000,
20            per_file_max_chars: 20_000,
21            total_context_max_chars: 100_000,
22            head_ratio: DEFAULT_HEAD_RATIO,
23            tail_ratio: DEFAULT_TAIL_RATIO,
24        }
25    }
26}
27
28pub fn truncate_with_marker(content: &str, max_chars: usize, name: &str) -> (String, bool) {
29    truncate_with_marker_and_ratio(
30        content,
31        max_chars,
32        name,
33        DEFAULT_HEAD_RATIO,
34        DEFAULT_TAIL_RATIO,
35    )
36}
37
38pub fn apply_budget(files: &mut Vec<ContextFile>, budget: &ContextBudget) {
39    for file in files.iter_mut() {
40        let (content, truncated) = truncate_with_marker_and_ratio(
41            &file.content,
42            budget.per_file_max_chars,
43            &file.name,
44            budget.head_ratio,
45            budget.tail_ratio,
46        );
47        file.content = content;
48        file.truncated |= truncated;
49    }
50
51    let mut prioritized = prioritized_files(files);
52    let mut remaining = budget.total_context_max_chars;
53    let mut kept = Vec::new();
54
55    for mut file in prioritized.drain(..) {
56        let file_chars = file.content.chars().count();
57        if file_chars <= remaining {
58            remaining -= file_chars;
59            kept.push(file);
60            continue;
61        }
62
63        let should_keep_as_truncated =
64            file.priority == ContextPriority::Critical || remaining >= MIN_FILE_ALLOCATION_CHARS;
65
66        if !should_keep_as_truncated {
67            continue;
68        }
69
70        if remaining == 0 {
71            continue;
72        }
73
74        let (content, truncated) = truncate_with_marker_and_ratio(
75            &file.content,
76            remaining,
77            &file.name,
78            budget.head_ratio,
79            budget.tail_ratio,
80        );
81        file.content = content;
82        file.truncated |= truncated;
83        remaining = 0;
84        kept.push(file);
85    }
86
87    *files = kept;
88}
89
90fn prioritized_files(files: &[ContextFile]) -> Vec<ContextFile> {
91    let mut prioritized = Vec::new();
92
93    for priority in [
94        ContextPriority::Critical,
95        ContextPriority::High,
96        ContextPriority::Normal,
97        ContextPriority::CallerSupplied,
98    ] {
99        for file in files {
100            if file.priority == priority {
101                prioritized.push(file.clone());
102            }
103        }
104    }
105
106    prioritized
107}
108
109fn truncate_with_marker_and_ratio(
110    content: &str,
111    max_chars: usize,
112    name: &str,
113    head_ratio: f64,
114    tail_ratio: f64,
115) -> (String, bool) {
116    if max_chars == 0 {
117        return (String::new(), !content.is_empty());
118    }
119
120    let chars: Vec<char> = content.chars().collect();
121    if chars.len() <= max_chars {
122        return (content.to_string(), false);
123    }
124
125    let marker = format!("\n[... truncated {name}; read file for full content ...]\n");
126    let marker_len = marker.chars().count();
127
128    if marker_len >= max_chars {
129        let truncated: String = chars.into_iter().take(max_chars).collect();
130        return (truncated, true);
131    }
132
133    let available = max_chars - marker_len;
134    let mut head_count = ((available as f64) * head_ratio).floor() as usize;
135    let mut tail_count = ((available as f64) * tail_ratio).floor() as usize;
136
137    if head_count + tail_count > available {
138        tail_count = tail_count.min(available.saturating_sub(head_count));
139    }
140
141    let used = head_count + tail_count;
142    if used < available {
143        head_count += available - used;
144    }
145
146    let head: String = chars.iter().take(head_count).collect();
147    let tail: String = chars
148        .iter()
149        .rev()
150        .take(tail_count)
151        .copied()
152        .collect::<Vec<char>>()
153        .into_iter()
154        .rev()
155        .collect();
156
157    (format!("{head}{marker}{tail}"), true)
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn truncates_with_marker() {
166        let content = "x".repeat(1_000);
167        let (truncated, changed) = truncate_with_marker(&content, 120, "AGENTS.md");
168
169        assert!(changed);
170        assert!(truncated.contains("truncated AGENTS.md"));
171        assert!(truncated.chars().count() <= 120);
172    }
173
174    #[test]
175    fn budget_prioritizes_critical_files() {
176        let mut files = vec![
177            ContextFile::new(
178                "notes",
179                "/tmp/notes",
180                "x".repeat(500),
181                ContextPriority::Normal,
182            ),
183            ContextFile::new(
184                "AGENTS.md",
185                "/tmp/AGENTS.md",
186                "y".repeat(500),
187                ContextPriority::Critical,
188            ),
189        ];
190
191        apply_budget(
192            &mut files,
193            &ContextBudget {
194                system_prompt_max_chars: 1_000,
195                per_file_max_chars: 1_000,
196                total_context_max_chars: 300,
197                head_ratio: 0.7,
198                tail_ratio: 0.2,
199            },
200        );
201
202        assert!(!files.is_empty());
203        assert_eq!(files[0].name, "AGENTS.md");
204    }
205}
stakpak_server/context/budget.rs

stakpak_server/context/
budget.rs
