tokenfold_core/
budget.rs

1use std::path::PathBuf;
2
3use crate::errors::TokenFoldError;
4use crate::input::{CompressionInput, InputFormat};
5use crate::token_estimator::TokenEstimator;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct CompressionPolicy {
9    pub target_tokens: Option<usize>,
10    pub reserve_output_tokens: usize,
11    pub mode: CompressionMode,
12    pub task_scope: TaskScope,
13    pub cache_boundary: Option<CacheBoundary>,
14    pub preserve_latest_user_message: bool,
15    pub disabled: Vec<String>,
16    pub unsafe_disable_redaction: bool,
17    /// CLI `--experimental`: enables transforms with `ModeEntry.experimental == true`
18    /// (currently `diff_compaction`; `log_compaction` was promoted out of `--experimental`
19    /// after the Phase 5 fidelity gate, see `modes::ALL_ENTRIES`) at their validated ratio band.
20    pub experimental: bool,
21    /// CLI `--enable <id>`: force-enable a specific transform ID even though its mode-matrix
22    /// entry doesn't enable it for the current mode. Still requires `experimental` for any
23    /// transform whose `ModeEntry.experimental == true` (see `modes::pipeline_for`).
24    pub enable: Vec<String>,
25    /// F-045: when true, and the full pre-transform input contains no secret-shaped content,
26    /// `pipeline::compress_with_estimator` persists it to the reversible evidence store
27    /// (`retrieval_backend`/`retrieval_store_path`) under its SHA-256 hash.
28    pub store_originals: bool,
29    /// F-045: the namespace stored-original entries are keyed under (see
30    /// `retrieval_store::RetrievalStore::store`).
31    pub retrieval_namespace: String,
32    /// F-045: TTL passed to `RetrievalStore::store` for newly stored originals. `None` means
33    /// "use `retrieval_store::DEFAULT_TTL_SECONDS`" (this is a *default*, not "never expire" —
34    /// that per-entry meaning belongs to `RetrievalStore::store`'s own `ttl_seconds` parameter).
35    pub retrieval_ttl_seconds: Option<u64>,
36    /// F-045: backend name passed to `RetrievalStore::open` ("memory" | "filesystem" |
37    /// "sqlite" — the latter fails clearly, handled as best-effort skip, see
38    /// `pipeline::maybe_store_originals`).
39    pub retrieval_backend: String,
40    /// F-045: filesystem backend root override. `None` means
41    /// `retrieval_store::default_store_path()`.
42    pub retrieval_store_path: Option<PathBuf>,
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum CompressionMode {
47    Conservative,
48    Balanced,
49    Aggressive,
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum TaskScope {
54    All,
55    General,
56    CodeReview,
57    ChangeSummary,
58    Debugging,
59    Generation,
60    ApiOverview,
61    RetrievalQa,
62    AgentHistory,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum CacheBoundary {
67    ByteOffset(usize),
68    TurnIndex(usize),
69}
70
71impl CompressionPolicy {
72    pub fn builder() -> CompressionPolicyBuilder {
73        CompressionPolicyBuilder::default()
74    }
75}
76
77#[derive(Debug, Clone, Default)]
78pub struct CompressionPolicyBuilder {
79    target_tokens: Option<usize>,
80    reserve_output_tokens: Option<usize>,
81    mode: Option<CompressionMode>,
82    task_scope: Option<TaskScope>,
83    cache_boundary: Option<CacheBoundary>,
84    preserve_latest_user_message: Option<bool>,
85    disabled: Vec<String>,
86    unsafe_disable_redaction: bool,
87    experimental: bool,
88    enable: Vec<String>,
89    store_originals: bool,
90    retrieval_namespace: Option<String>,
91    retrieval_ttl_seconds: Option<u64>,
92    retrieval_backend: Option<String>,
93    retrieval_store_path: Option<PathBuf>,
94}
95
96impl CompressionPolicyBuilder {
97    pub fn target_tokens(mut self, target_tokens: usize) -> Self {
98        self.target_tokens = Some(target_tokens);
99        self
100    }
101
102    pub fn reserve_output_tokens(mut self, reserve_output_tokens: usize) -> Self {
103        self.reserve_output_tokens = Some(reserve_output_tokens);
104        self
105    }
106
107    pub fn mode(mut self, mode: CompressionMode) -> Self {
108        self.mode = Some(mode);
109        self
110    }
111
112    pub fn task_scope(mut self, task_scope: TaskScope) -> Self {
113        self.task_scope = Some(task_scope);
114        self
115    }
116
117    pub fn cache_boundary(mut self, cache_boundary: CacheBoundary) -> Self {
118        self.cache_boundary = Some(cache_boundary);
119        self
120    }
121
122    pub fn preserve_latest_user_message(mut self, preserve: bool) -> Self {
123        self.preserve_latest_user_message = Some(preserve);
124        self
125    }
126
127    pub fn disable(mut self, transform_id: impl Into<String>) -> Self {
128        self.disabled.push(transform_id.into());
129        self
130    }
131
132    pub fn unsafe_disable_redaction(mut self, unsafe_disable: bool) -> Self {
133        self.unsafe_disable_redaction = unsafe_disable;
134        self
135    }
136
137    pub fn experimental(mut self, experimental: bool) -> Self {
138        self.experimental = experimental;
139        self
140    }
141
142    pub fn enable(mut self, transform_id: impl Into<String>) -> Self {
143        self.enable.push(transform_id.into());
144        self
145    }
146
147    pub fn store_originals(mut self, store_originals: bool) -> Self {
148        self.store_originals = store_originals;
149        self
150    }
151
152    pub fn retrieval_namespace(mut self, namespace: impl Into<String>) -> Self {
153        self.retrieval_namespace = Some(namespace.into());
154        self
155    }
156
157    pub fn retrieval_ttl_seconds(mut self, ttl_seconds: Option<u64>) -> Self {
158        self.retrieval_ttl_seconds = ttl_seconds;
159        self
160    }
161
162    pub fn retrieval_backend(mut self, backend: impl Into<String>) -> Self {
163        self.retrieval_backend = Some(backend.into());
164        self
165    }
166
167    pub fn retrieval_store_path(mut self, store_path: Option<PathBuf>) -> Self {
168        self.retrieval_store_path = store_path;
169        self
170    }
171
172    pub fn build(self) -> Result<CompressionPolicy, TokenFoldError> {
173        if self.disabled.iter().any(|id| id == "secret_redaction") {
174            return Err(TokenFoldError::ConfigError(
175                "secret_redaction cannot be disabled via CompressionPolicy.disabled".to_string(),
176            ));
177        }
178        Ok(CompressionPolicy {
179            target_tokens: self.target_tokens,
180            reserve_output_tokens: self.reserve_output_tokens.unwrap_or(0),
181            mode: self.mode.unwrap_or(CompressionMode::Balanced),
182            task_scope: self.task_scope.unwrap_or(TaskScope::All),
183            cache_boundary: self.cache_boundary,
184            preserve_latest_user_message: self.preserve_latest_user_message.unwrap_or(true),
185            disabled: self.disabled,
186            unsafe_disable_redaction: self.unsafe_disable_redaction,
187            experimental: self.experimental,
188            enable: self.enable,
189            store_originals: self.store_originals,
190            retrieval_namespace: self
191                .retrieval_namespace
192                .unwrap_or_else(|| "default".to_string()),
193            retrieval_ttl_seconds: self.retrieval_ttl_seconds,
194            retrieval_backend: self
195                .retrieval_backend
196                .unwrap_or_else(|| "filesystem".to_string()),
197            retrieval_store_path: self.retrieval_store_path,
198        })
199    }
200}
201
202/// tokens(protected + structurally-required content). Used to detect `Status::UnreachableTarget`.
203pub fn protected_floor(
204    input: &CompressionInput,
205    policy: &CompressionPolicy,
206    estimator: &dyn TokenEstimator,
207) -> usize {
208    estimator.count_bytes(&protected_segments(input, policy).concat())
209}
210
211/// The individual protected-content segments (one per system message, the latest user
212/// message, each diff header/hunk line, …) that must each survive byte-for-byte after any
213/// transform. Kept as separate segments (rather than one flattened blob) so `safety.rs` can
214/// check each one independently — concatenated messages are rarely contiguous in the
215/// original document, so a single substring check across the whole blob would be meaningless.
216pub fn protected_segments(input: &CompressionInput, policy: &CompressionPolicy) -> Vec<Vec<u8>> {
217    match input.format {
218        InputFormat::OpenAiJson => extract_openai_protected(&input.bytes, policy),
219        InputFormat::AnthropicJson => extract_anthropic_protected(&input.bytes, policy),
220        InputFormat::GitDiff => extract_diff_protected(&input.bytes),
221        // ponytail: no transform touches plain text/command output structure yet beyond
222        // log/diff compaction (task-scope gated), so nothing is unconditionally protected.
223        InputFormat::PlainText | InputFormat::CommandOutput | InputFormat::Auto => Vec::new(),
224    }
225}
226
227fn extract_openai_protected(bytes: &[u8], policy: &CompressionPolicy) -> Vec<Vec<u8>> {
228    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
229        return Vec::new();
230    };
231    let Some(messages) = value.get("messages").and_then(|m| m.as_array()) else {
232        return Vec::new();
233    };
234
235    let mut segments = Vec::new();
236    for message in messages {
237        if message.get("role").and_then(|r| r.as_str()) == Some("system")
238            && let Some(bytes) = message_content_bytes(message)
239        {
240            segments.push(bytes);
241        }
242    }
243    if policy.preserve_latest_user_message
244        && let Some(last_user) = messages
245            .iter()
246            .rev()
247            .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
248        && let Some(bytes) = message_content_bytes(last_user)
249    {
250        segments.push(bytes);
251    }
252    segments
253}
254
255fn extract_anthropic_protected(bytes: &[u8], policy: &CompressionPolicy) -> Vec<Vec<u8>> {
256    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
257        return Vec::new();
258    };
259
260    let mut segments = Vec::new();
261    if let Some(system) = value.get("system").and_then(|s| s.as_str()) {
262        segments.push(system.as_bytes().to_vec());
263    }
264    if policy.preserve_latest_user_message
265        && let Some(last_user) =
266            value
267                .get("messages")
268                .and_then(|m| m.as_array())
269                .and_then(|messages| {
270                    messages
271                        .iter()
272                        .rev()
273                        .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
274                })
275        && let Some(bytes) = message_content_bytes(last_user)
276    {
277        segments.push(bytes);
278    }
279    segments
280}
281
282fn message_content_bytes(message: &serde_json::Value) -> Option<Vec<u8>> {
283    match message.get("content") {
284        Some(serde_json::Value::String(text)) => Some(text.as_bytes().to_vec()),
285        Some(structured) => serde_json::to_vec(structured).ok(),
286        None => None,
287    }
288}
289
290/// Keeps file names and hunk headers, matching the `diff_compaction` (F-013) contract of what
291/// must survive compaction. Each kept line is its own segment.
292fn extract_diff_protected(bytes: &[u8]) -> Vec<Vec<u8>> {
293    let text = String::from_utf8_lossy(bytes);
294    let mut segments = Vec::new();
295    for line in text.lines() {
296        if line.starts_with("diff --git")
297            || line.starts_with("--- ")
298            || line.starts_with("+++ ")
299            || line.starts_with("@@")
300        {
301            let mut segment = line.as_bytes().to_vec();
302            segment.push(b'\n');
303            segments.push(segment);
304        }
305    }
306    segments
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312    use crate::token_estimator::ByteHeuristicEstimator;
313
314    #[test]
315    fn default_mode_is_balanced() {
316        let policy = CompressionPolicy::builder().build().unwrap();
317        assert_eq!(policy.mode, CompressionMode::Balanced);
318    }
319
320    #[test]
321    fn store_originals_defaults_to_false_with_a_default_namespace() {
322        let policy = CompressionPolicy::builder().build().unwrap();
323        assert!(!policy.store_originals);
324        assert_eq!(policy.retrieval_namespace, "default");
325    }
326
327    #[test]
328    fn store_originals_and_namespace_are_settable_via_the_builder() {
329        let policy = CompressionPolicy::builder()
330            .store_originals(true)
331            .retrieval_namespace("project-x")
332            .retrieval_ttl_seconds(Some(60))
333            .retrieval_backend("memory")
334            .retrieval_store_path(Some(std::path::PathBuf::from("/tmp/custom")))
335            .build()
336            .unwrap();
337        assert!(policy.store_originals);
338        assert_eq!(policy.retrieval_namespace, "project-x");
339        assert_eq!(policy.retrieval_ttl_seconds, Some(60));
340        assert_eq!(policy.retrieval_backend, "memory");
341        assert_eq!(
342            policy.retrieval_store_path,
343            Some(std::path::PathBuf::from("/tmp/custom"))
344        );
345    }
346
347    #[test]
348    fn retrieval_defaults_are_none_ttl_and_filesystem_backend() {
349        let policy = CompressionPolicy::builder().build().unwrap();
350        assert_eq!(policy.retrieval_ttl_seconds, None);
351        assert_eq!(policy.retrieval_backend, "filesystem");
352        assert_eq!(policy.retrieval_store_path, None);
353    }
354
355    #[test]
356    fn secret_redaction_cannot_be_disabled_through_policy() {
357        let err = CompressionPolicy::builder()
358            .disable("secret_redaction")
359            .build()
360            .unwrap_err();
361        assert!(matches!(err, TokenFoldError::ConfigError(_)));
362    }
363
364    #[test]
365    fn disabling_other_transforms_is_allowed() {
366        let policy = CompressionPolicy::builder()
367            .disable("json_minify")
368            .build()
369            .unwrap();
370        assert_eq!(policy.disabled, vec!["json_minify".to_string()]);
371    }
372
373    #[test]
374    fn floor_is_zero_for_plain_text() {
375        let input = CompressionInput::plain_text(b"just some plain text".to_vec());
376        let policy = CompressionPolicy::builder().build().unwrap();
377        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
378        assert_eq!(floor, 0);
379    }
380
381    #[test]
382    fn floor_covers_system_and_latest_user_message_for_openai_json() {
383        let payload = serde_json::json!({
384            "model": "gpt-4",
385            "messages": [
386                {"role": "system", "content": "You are a helpful assistant."},
387                {"role": "user", "content": "first question"},
388                {"role": "assistant", "content": "first answer"},
389                {"role": "user", "content": "second question"},
390            ]
391        });
392        let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
393        let policy = CompressionPolicy::builder().build().unwrap();
394        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
395
396        let expected_bytes = "You are a helpful assistant.".len() + "second question".len();
397        assert_eq!(
398            floor,
399            ByteHeuristicEstimator.count_bytes(&vec![0u8; expected_bytes])
400        );
401        // The earlier "first question" turn must NOT be counted as protected.
402        assert!(floor < ByteHeuristicEstimator.count_bytes(input.bytes.as_slice()));
403    }
404
405    #[test]
406    fn floor_excludes_latest_user_message_when_policy_disables_preservation() {
407        let payload = serde_json::json!({
408            "messages": [
409                {"role": "system", "content": "system prompt"},
410                {"role": "user", "content": "question"},
411            ]
412        });
413        let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
414        let policy = CompressionPolicy::builder()
415            .preserve_latest_user_message(false)
416            .build()
417            .unwrap();
418        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
419        assert_eq!(floor, ByteHeuristicEstimator.count_bytes(b"system prompt"));
420    }
421
422    #[test]
423    fn floor_covers_system_and_latest_user_message_for_anthropic_json() {
424        let payload = serde_json::json!({
425            "system": "system prompt",
426            "messages": [
427                {"role": "user", "content": "first"},
428                {"role": "assistant", "content": "reply"},
429                {"role": "user", "content": "second"},
430            ]
431        });
432        let input = CompressionInput::anthropic_json(serde_json::to_vec(&payload).unwrap());
433        let policy = CompressionPolicy::builder().build().unwrap();
434        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
435        let expected_bytes = "system prompt".len() + "second".len();
436        assert_eq!(
437            floor,
438            ByteHeuristicEstimator.count_bytes(&vec![0u8; expected_bytes])
439        );
440    }
441
442    #[test]
443    fn floor_keeps_diff_headers_and_hunk_markers_only() {
444        let diff =
445            b"diff --git a/f.rs b/f.rs\n--- a/f.rs\n+++ b/f.rs\n@@ -1,2 +1,2 @@\n-old\n+new\n";
446        let input = CompressionInput::git_diff(diff.to_vec());
447        let policy = CompressionPolicy::builder().build().unwrap();
448        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
449        assert!(floor > 0);
450        assert!(floor < ByteHeuristicEstimator.count_bytes(diff));
451    }
452
453    #[test]
454    fn malformed_json_never_panics_and_yields_zero_floor() {
455        let input = CompressionInput::openai_json(b"{not json".to_vec());
456        let policy = CompressionPolicy::builder().build().unwrap();
457        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
458        assert_eq!(floor, 0);
459    }
460}
tokenfold_core/budget.rs

tokenfold_core/
budget.rs
