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        // Generic Json has no "protected" sub-segment either — json_field_fold's own
224        // round-trip safety gate is what guarantees its data is preserved.
225        InputFormat::PlainText
226        | InputFormat::CommandOutput
227        | InputFormat::Json
228        | InputFormat::Auto => Vec::new(),
229    }
230}
231
232fn extract_openai_protected(bytes: &[u8], policy: &CompressionPolicy) -> Vec<Vec<u8>> {
233    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
234        return Vec::new();
235    };
236    let Some(messages) = value.get("messages").and_then(|m| m.as_array()) else {
237        return Vec::new();
238    };
239
240    let mut segments = Vec::new();
241    for message in messages {
242        if message.get("role").and_then(|r| r.as_str()) == Some("system")
243            && let Some(bytes) = message_content_bytes(message)
244        {
245            segments.push(bytes);
246        }
247    }
248    if policy.preserve_latest_user_message
249        && let Some(last_user) = messages
250            .iter()
251            .rev()
252            .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
253        && let Some(bytes) = message_content_bytes(last_user)
254    {
255        segments.push(bytes);
256    }
257    segments
258}
259
260fn extract_anthropic_protected(bytes: &[u8], policy: &CompressionPolicy) -> Vec<Vec<u8>> {
261    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
262        return Vec::new();
263    };
264
265    let mut segments = Vec::new();
266    if let Some(system) = value.get("system").and_then(|s| s.as_str()) {
267        segments.push(system.as_bytes().to_vec());
268    }
269    if policy.preserve_latest_user_message
270        && let Some(last_user) =
271            value
272                .get("messages")
273                .and_then(|m| m.as_array())
274                .and_then(|messages| {
275                    messages
276                        .iter()
277                        .rev()
278                        .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
279                })
280        && let Some(bytes) = message_content_bytes(last_user)
281    {
282        segments.push(bytes);
283    }
284    segments
285}
286
287fn message_content_bytes(message: &serde_json::Value) -> Option<Vec<u8>> {
288    match message.get("content") {
289        Some(serde_json::Value::String(text)) => Some(text.as_bytes().to_vec()),
290        Some(structured) => serde_json::to_vec(structured).ok(),
291        None => None,
292    }
293}
294
295/// Keeps file names and hunk headers, matching the `diff_compaction` (F-013) contract of what
296/// must survive compaction. Each kept line is its own segment.
297fn extract_diff_protected(bytes: &[u8]) -> Vec<Vec<u8>> {
298    let text = String::from_utf8_lossy(bytes);
299    let mut segments = Vec::new();
300    for line in text.lines() {
301        if line.starts_with("diff --git")
302            || line.starts_with("--- ")
303            || line.starts_with("+++ ")
304            || line.starts_with("@@")
305        {
306            let mut segment = line.as_bytes().to_vec();
307            segment.push(b'\n');
308            segments.push(segment);
309        }
310    }
311    segments
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317    use crate::token_estimator::ByteHeuristicEstimator;
318
319    #[test]
320    fn default_mode_is_balanced() {
321        let policy = CompressionPolicy::builder().build().unwrap();
322        assert_eq!(policy.mode, CompressionMode::Balanced);
323    }
324
325    #[test]
326    fn store_originals_defaults_to_false_with_a_default_namespace() {
327        let policy = CompressionPolicy::builder().build().unwrap();
328        assert!(!policy.store_originals);
329        assert_eq!(policy.retrieval_namespace, "default");
330    }
331
332    #[test]
333    fn store_originals_and_namespace_are_settable_via_the_builder() {
334        let policy = CompressionPolicy::builder()
335            .store_originals(true)
336            .retrieval_namespace("project-x")
337            .retrieval_ttl_seconds(Some(60))
338            .retrieval_backend("memory")
339            .retrieval_store_path(Some(std::path::PathBuf::from("/tmp/custom")))
340            .build()
341            .unwrap();
342        assert!(policy.store_originals);
343        assert_eq!(policy.retrieval_namespace, "project-x");
344        assert_eq!(policy.retrieval_ttl_seconds, Some(60));
345        assert_eq!(policy.retrieval_backend, "memory");
346        assert_eq!(
347            policy.retrieval_store_path,
348            Some(std::path::PathBuf::from("/tmp/custom"))
349        );
350    }
351
352    #[test]
353    fn retrieval_defaults_are_none_ttl_and_filesystem_backend() {
354        let policy = CompressionPolicy::builder().build().unwrap();
355        assert_eq!(policy.retrieval_ttl_seconds, None);
356        assert_eq!(policy.retrieval_backend, "filesystem");
357        assert_eq!(policy.retrieval_store_path, None);
358    }
359
360    #[test]
361    fn secret_redaction_cannot_be_disabled_through_policy() {
362        let err = CompressionPolicy::builder()
363            .disable("secret_redaction")
364            .build()
365            .unwrap_err();
366        assert!(matches!(err, TokenFoldError::ConfigError(_)));
367    }
368
369    #[test]
370    fn disabling_other_transforms_is_allowed() {
371        let policy = CompressionPolicy::builder()
372            .disable("json_minify")
373            .build()
374            .unwrap();
375        assert_eq!(policy.disabled, vec!["json_minify".to_string()]);
376    }
377
378    #[test]
379    fn floor_is_zero_for_plain_text() {
380        let input = CompressionInput::plain_text(b"just some plain text".to_vec());
381        let policy = CompressionPolicy::builder().build().unwrap();
382        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
383        assert_eq!(floor, 0);
384    }
385
386    #[test]
387    fn floor_covers_system_and_latest_user_message_for_openai_json() {
388        let payload = serde_json::json!({
389            "model": "gpt-4",
390            "messages": [
391                {"role": "system", "content": "You are a helpful assistant."},
392                {"role": "user", "content": "first question"},
393                {"role": "assistant", "content": "first answer"},
394                {"role": "user", "content": "second question"},
395            ]
396        });
397        let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
398        let policy = CompressionPolicy::builder().build().unwrap();
399        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
400
401        let expected_bytes = "You are a helpful assistant.".len() + "second question".len();
402        assert_eq!(
403            floor,
404            ByteHeuristicEstimator.count_bytes(&vec![0u8; expected_bytes])
405        );
406        // The earlier "first question" turn must NOT be counted as protected.
407        assert!(floor < ByteHeuristicEstimator.count_bytes(input.bytes.as_slice()));
408    }
409
410    #[test]
411    fn floor_excludes_latest_user_message_when_policy_disables_preservation() {
412        let payload = serde_json::json!({
413            "messages": [
414                {"role": "system", "content": "system prompt"},
415                {"role": "user", "content": "question"},
416            ]
417        });
418        let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
419        let policy = CompressionPolicy::builder()
420            .preserve_latest_user_message(false)
421            .build()
422            .unwrap();
423        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
424        assert_eq!(floor, ByteHeuristicEstimator.count_bytes(b"system prompt"));
425    }
426
427    #[test]
428    fn floor_covers_system_and_latest_user_message_for_anthropic_json() {
429        let payload = serde_json::json!({
430            "system": "system prompt",
431            "messages": [
432                {"role": "user", "content": "first"},
433                {"role": "assistant", "content": "reply"},
434                {"role": "user", "content": "second"},
435            ]
436        });
437        let input = CompressionInput::anthropic_json(serde_json::to_vec(&payload).unwrap());
438        let policy = CompressionPolicy::builder().build().unwrap();
439        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
440        let expected_bytes = "system prompt".len() + "second".len();
441        assert_eq!(
442            floor,
443            ByteHeuristicEstimator.count_bytes(&vec![0u8; expected_bytes])
444        );
445    }
446
447    #[test]
448    fn floor_keeps_diff_headers_and_hunk_markers_only() {
449        let diff =
450            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";
451        let input = CompressionInput::git_diff(diff.to_vec());
452        let policy = CompressionPolicy::builder().build().unwrap();
453        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
454        assert!(floor > 0);
455        assert!(floor < ByteHeuristicEstimator.count_bytes(diff));
456    }
457
458    #[test]
459    fn malformed_json_never_panics_and_yields_zero_floor() {
460        let input = CompressionInput::openai_json(b"{not json".to_vec());
461        let policy = CompressionPolicy::builder().build().unwrap();
462        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
463        assert_eq!(floor, 0);
464    }
465}
tokenfold_core/budget.rs

tokenfold_core/
budget.rs
