tokenfold_core/
budget.rs

1use std::path::PathBuf;
2
3use crate::codec::OutputEncoding;
4use crate::errors::TokenFoldError;
5use crate::input::{CompressionInput, InputFormat};
6use crate::token_estimator::TokenEstimator;
7
8/// Placeholder floor for `retrieval_ttl_seconds` whenever `lossy` is set — see
9/// `CompressionPolicyBuilder::build`'s lossy validation. The exact retention policy is still an
10/// open owner decision, not yet confirmed; this value only needs to be "clearly more than
11/// instant," not final.
12const MIN_LOSSY_TTL_SECONDS: u64 = 86_400;
13
14// NOT Eq: `lossy_ratio` holds an f64 (same precedent as `CompressionOutput`/`CompressionReport`).
15#[derive(Debug, Clone, PartialEq)]
16pub struct CompressionPolicy {
17    pub target_tokens: Option<usize>,
18    pub reserve_output_tokens: usize,
19    pub preset: Preset,
20    pub task_scope: TaskScope,
21    pub encoding: OutputEncoding,
22    pub pruning: Option<PruningPolicy>,
23    pub preserve_latest_user_message: bool,
24    pub disabled: Vec<String>,
25    /// CLI `--experimental`: enables transforms with `ModeEntry.experimental == true`
26    /// (currently `diff_compaction`; `log_compaction` was promoted out of `--experimental`
27    /// after the Phase 5 fidelity gate, see `modes::ALL_ENTRIES`) at their validated ratio band.
28    pub experimental: bool,
29    /// CLI `--enable <id>`: force-enable a specific transform ID even though its preset-matrix
30    /// entry doesn't enable it for the current preset. Still requires `experimental` for any
31    /// transform whose `ModeEntry.experimental == true` (see `modes::pipeline_for`).
32    pub enable: Vec<String>,
33    /// When true, and the full pre-transform input contains no secret-shaped content,
34    /// `pipeline::compress_with_estimator` persists it to the reversible evidence store
35    /// (`retrieval_backend`/`retrieval_store_path`) under its SHA-256 hash.
36    pub store_originals: bool,
37    /// The evidence-store namespace stored-original entries are keyed under (see
38    /// `retrieval_store::RetrievalStore::store`).
39    pub retrieval_namespace: String,
40    /// TTL passed to `RetrievalStore::store` for newly stored originals. `None` means
41    /// "use `retrieval_store::DEFAULT_TTL_SECONDS`" (this is a *default*, not "never expire" —
42    /// that per-entry meaning belongs to `RetrievalStore::store`'s own `ttl_seconds` parameter).
43    pub retrieval_ttl_seconds: Option<u64>,
44    /// Backend name passed to `RetrievalStore::open` ("memory" | "filesystem" |
45    /// "sqlite" — the latter fails clearly, handled as best-effort skip, see
46    /// `pipeline::maybe_store_originals`).
47    pub retrieval_backend: String,
48    /// Filesystem backend root override. `None` means
49    /// `retrieval_store::default_store_path()`.
50    pub retrieval_store_path: Option<PathBuf>,
51    /// Opt-in lossy JSON array-item selection: array items are dropped and replaced by a
52    /// recoverable `$tf_ref` marker. `None` (the default) means the lossless pipeline is
53    /// untouched — this field is set only by an explicit CLI flag, never derived from
54    /// `preset`/`experimental`, and is deliberately NOT part of `modes.rs`/`ALL_ENTRIES`: it is a
55    /// fundamentally different category (data-lossy, not just structurally-lossy-but-reversible)
56    /// from every other transform in this crate. See `pipeline::apply_lossy_reduction`.
57    pub lossy: Option<LossyPath>,
58    /// BEST-EFFORT selection hint, not an enforced budget: how aggressively to prune, as the
59    /// fraction (0.0..=1.0) of the prunable pool's own estimated token cost to keep when `lossy`
60    /// is set. It parameterizes `transforms::json_prune`'s selection walk and is deliberately
61    /// never re-checked against the final serialized document — the achieved whole-document ratio
62    /// will differ, since the pool excludes preserved arrays, all non-array content, and items
63    /// cheaper than the `$tf_ref` marker that would replace them, and since
64    /// `pipeline::apply_lossy_reduction` discards a whole prune that fails to beat the lossless
65    /// pipeline. `target_tokens` is the enforced ceiling; this is not. Ignored when `lossy` is
66    /// `None`.
67    pub lossy_ratio: f64,
68    /// Dot-separated paths (see `transforms::json_prune::LossyOptions::preserve_paths`) whose
69    /// arrays must never be pruned. Ignored when `lossy` is `None`.
70    pub lossy_preserve: Vec<String>,
71    /// True for a side-effect-free preview (`tokenfold inspect` / `compress --dry-run`): the
72    /// projected output/savings are computed exactly as a real run would, but
73    /// `pipeline::maybe_store_originals`/`apply_lossy_reduction` must not perform any real
74    /// `RetrievalStore` write.
75    ///
76    /// `pub(crate)`, settable only via [`CompressionPolicyBuilder::preview`]. It was a plain
77    /// `pub` field, which made it a footgun for library callers: a preview run's output can carry
78    /// `$tf_ref` markers whose targets were deliberately never persisted, so flipping this on an
79    /// otherwise ordinary policy and feeding `CompressionOutput::bytes` to a model yields
80    /// references that resolve to nothing. Callers who want the projection must ask for it by
81    /// name and are told, right here, that the bytes are a projection to measure — not to ship.
82    pub(crate) preview: bool,
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum Preset {
87    Conservative,
88    Balanced,
89    Aggressive,
90}
91
92/// Selection backend for opt-in lossy JSON pruning. `Heuristic` is the only Phase 1
93/// implementation; a future `Select` (Tokenfold Select as the scorer) is Phase 2 and not
94/// implemented — deliberately a single-variant enum for now rather than a bare `bool`, since the
95/// CLI already speaks of this as choosing an algorithm (`--lossy heuristic`), not toggling a flag.
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub enum LossyPath {
98    Heuristic,
99}
100
101#[derive(Debug, Clone, PartialEq)]
102pub struct PruningPolicy {
103    pub keep_ratio: Option<f64>,
104    pub preserve_paths: Vec<String>,
105    pub retrieval_store: Option<PathBuf>,
106    pub retrieval_namespace: Option<String>,
107}
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub enum TaskScope {
111    All,
112    General,
113    CodeReview,
114    ChangeSummary,
115    Debugging,
116    Generation,
117    ApiOverview,
118    RetrievalQa,
119    AgentHistory,
120}
121
122impl CompressionPolicy {
123    pub fn builder() -> CompressionPolicyBuilder {
124        CompressionPolicyBuilder::default()
125    }
126
127    /// Re-checks the same invariants `CompressionPolicyBuilder::build` enforces, against the
128    /// concrete built struct rather than the builder's `Option` fields. Every field here is
129    /// `pub`, so a caller can construct or mutate a `CompressionPolicy` directly without ever
130    /// going through the builder -- `pipeline::compress_with_estimator` calls this on every
131    /// policy it receives so a hand-built policy can't silently skip the same fail-closed
132    /// guarantees a builder-built one gets for free.
133    pub fn validate(&self) -> Result<(), TokenFoldError> {
134        if let Some(pruning) = &self.pruning {
135            if pruning.keep_ratio.is_none() && self.target_tokens.is_none() {
136                return Err(TokenFoldError::ConfigError(
137                    "pruning requires target_tokens or keep_ratio".to_string(),
138                ));
139            }
140            if pruning
141                .keep_ratio
142                .is_some_and(|ratio| !(0.0 < ratio && ratio <= 1.0))
143            {
144                return Err(TokenFoldError::ConfigError(
145                    "keep_ratio must be greater than 0 and at most 1".to_string(),
146                ));
147            }
148        }
149        if self.disabled.iter().any(|id| id == "secret_redaction") {
150            return Err(TokenFoldError::ConfigError(
151                "secret_redaction cannot be disabled via CompressionPolicy.disabled".to_string(),
152            ));
153        }
154        if !(0.0..=1.0).contains(&self.lossy_ratio) {
155            return Err(TokenFoldError::ConfigError(format!(
156                "lossy_ratio must be between 0.0 and 1.0, got {}",
157                self.lossy_ratio
158            )));
159        }
160        if self.lossy.is_some() {
161            // Design doc §4/§8: "must refuse to run when retrieval_backend == Memory or the TTL
162            // is below a floor" -- a lossy run with no durable receipt is real data loss, not
163            // "lossy but recoverable".
164            if self.retrieval_backend != "filesystem" {
165                return Err(TokenFoldError::ConfigError(format!(
166                    "lossy pruning requires a durable retrieval backend (\"filesystem\"); \
167                     {:?} would make dropped items unrecoverable",
168                    self.retrieval_backend
169                )));
170            }
171            let effective_ttl = self
172                .retrieval_ttl_seconds
173                .unwrap_or(crate::retrieval_store::DEFAULT_TTL_SECONDS);
174            if effective_ttl < MIN_LOSSY_TTL_SECONDS {
175                return Err(TokenFoldError::ConfigError(format!(
176                    "lossy pruning requires retrieval_ttl_seconds >= {MIN_LOSSY_TTL_SECONDS} \
177                     (got {effective_ttl}); a near-immediate expiry has no real recoverability"
178                )));
179            }
180        }
181        Ok(())
182    }
183}
184
185#[derive(Debug, Clone, Default)]
186pub struct CompressionPolicyBuilder {
187    target_tokens: Option<usize>,
188    reserve_output_tokens: Option<usize>,
189    preset: Option<Preset>,
190    task_scope: Option<TaskScope>,
191    encoding: Option<OutputEncoding>,
192    pruning: Option<PruningPolicy>,
193    preserve_latest_user_message: Option<bool>,
194    disabled: Vec<String>,
195    experimental: bool,
196    enable: Vec<String>,
197    store_originals: bool,
198    retrieval_namespace: Option<String>,
199    retrieval_ttl_seconds: Option<u64>,
200    retrieval_backend: Option<String>,
201    retrieval_store_path: Option<PathBuf>,
202    lossy: Option<LossyPath>,
203    lossy_ratio: Option<f64>,
204    lossy_preserve: Vec<String>,
205    preview: bool,
206}
207
208impl CompressionPolicyBuilder {
209    pub fn target_tokens(mut self, target_tokens: usize) -> Self {
210        self.target_tokens = Some(target_tokens);
211        self
212    }
213
214    pub fn reserve_output_tokens(mut self, reserve_output_tokens: usize) -> Self {
215        self.reserve_output_tokens = Some(reserve_output_tokens);
216        self
217    }
218
219    pub fn preset(mut self, preset: Preset) -> Self {
220        self.preset = Some(preset);
221        self
222    }
223
224    pub fn task_scope(mut self, task_scope: TaskScope) -> Self {
225        self.task_scope = Some(task_scope);
226        self
227    }
228
229    pub fn encoding(mut self, encoding: OutputEncoding) -> Self {
230        self.encoding = Some(encoding);
231        self
232    }
233
234    pub fn pruning(mut self, pruning: PruningPolicy) -> Self {
235        self.lossy = Some(LossyPath::Heuristic);
236        self.lossy_ratio = Some(pruning.keep_ratio.unwrap_or(0.0));
237        self.lossy_preserve = pruning.preserve_paths.clone();
238        if pruning.retrieval_store.is_some() {
239            self.retrieval_store_path = pruning.retrieval_store.clone();
240        }
241        if pruning.retrieval_namespace.is_some() {
242            self.retrieval_namespace = pruning.retrieval_namespace.clone();
243        }
244        self.pruning = Some(pruning);
245        self
246    }
247
248    pub fn preserve_latest_user_message(mut self, preserve: bool) -> Self {
249        self.preserve_latest_user_message = Some(preserve);
250        self
251    }
252
253    pub fn disable(mut self, transform_id: impl Into<String>) -> Self {
254        self.disabled.push(transform_id.into());
255        self
256    }
257
258    pub fn experimental(mut self, experimental: bool) -> Self {
259        self.experimental = experimental;
260        self
261    }
262
263    pub fn enable(mut self, transform_id: impl Into<String>) -> Self {
264        self.enable.push(transform_id.into());
265        self
266    }
267
268    pub fn store_originals(mut self, store_originals: bool) -> Self {
269        self.store_originals = store_originals;
270        self
271    }
272
273    pub fn retrieval_namespace(mut self, namespace: impl Into<String>) -> Self {
274        self.retrieval_namespace = Some(namespace.into());
275        self
276    }
277
278    pub fn retrieval_ttl_seconds(mut self, ttl_seconds: Option<u64>) -> Self {
279        self.retrieval_ttl_seconds = ttl_seconds;
280        self
281    }
282
283    pub fn retrieval_backend(mut self, backend: impl Into<String>) -> Self {
284        self.retrieval_backend = Some(backend.into());
285        self
286    }
287
288    pub fn retrieval_store_path(mut self, store_path: Option<PathBuf>) -> Self {
289        self.retrieval_store_path = store_path;
290        self
291    }
292
293    pub fn lossy(mut self, lossy: LossyPath) -> Self {
294        self.lossy = Some(lossy);
295        self
296    }
297
298    pub fn lossy_ratio(mut self, ratio: f64) -> Self {
299        self.lossy_ratio = Some(ratio);
300        self
301    }
302
303    pub fn lossy_preserve(mut self, path: impl Into<String>) -> Self {
304        self.lossy_preserve.push(path.into());
305        self
306    }
307
308    /// Opt into a side-effect-free preview: no real `RetrievalStore` write happens anywhere in
309    /// the pipeline. The returned `CompressionOutput::bytes` are a PROJECTION of what a real run
310    /// would emit — with a lossy policy they can contain `$tf_ref` markers pointing at content
311    /// that was deliberately never stored, so they are for measuring savings, never for sending
312    /// to a model. See `CompressionPolicy::preview`.
313    pub fn preview(mut self, preview: bool) -> Self {
314        self.preview = preview;
315        self
316    }
317
318    pub fn build(self) -> Result<CompressionPolicy, TokenFoldError> {
319        let policy = CompressionPolicy {
320            target_tokens: self.target_tokens,
321            reserve_output_tokens: self.reserve_output_tokens.unwrap_or(0),
322            preset: self.preset.unwrap_or(Preset::Balanced),
323            task_scope: self.task_scope.unwrap_or(TaskScope::All),
324            encoding: self.encoding.unwrap_or_default(),
325            pruning: self.pruning,
326            preserve_latest_user_message: self.preserve_latest_user_message.unwrap_or(true),
327            disabled: self.disabled,
328            experimental: self.experimental,
329            enable: self.enable,
330            store_originals: self.store_originals,
331            retrieval_namespace: self
332                .retrieval_namespace
333                .unwrap_or_else(|| "default".to_string()),
334            retrieval_ttl_seconds: self.retrieval_ttl_seconds,
335            retrieval_backend: self
336                .retrieval_backend
337                .unwrap_or_else(|| "filesystem".to_string()),
338            retrieval_store_path: self.retrieval_store_path,
339            lossy: self.lossy,
340            lossy_ratio: self.lossy_ratio.unwrap_or(0.3),
341            lossy_preserve: self.lossy_preserve,
342            preview: self.preview,
343        };
344        policy.validate()?;
345        Ok(policy)
346    }
347}
348
349/// tokens(protected + structurally-required content). Used to detect `Status::UnreachableTarget`.
350pub fn protected_floor(
351    input: &CompressionInput,
352    policy: &CompressionPolicy,
353    estimator: &dyn TokenEstimator,
354) -> usize {
355    estimator.count_bytes(&protected_segments(input, policy).concat())
356}
357
358/// The individual protected-content segments (one per system message, the latest user
359/// message, each diff header/hunk line, …) that must each survive byte-for-byte after any
360/// transform. Kept as separate segments (rather than one flattened blob) so `safety.rs` can
361/// check each one independently — concatenated messages are rarely contiguous in the
362/// original document, so a single substring check across the whole blob would be meaningless.
363pub fn protected_segments(input: &CompressionInput, policy: &CompressionPolicy) -> Vec<Vec<u8>> {
364    match input.format {
365        InputFormat::OpenAiJson => extract_openai_protected(&input.bytes, policy),
366        InputFormat::AnthropicJson => extract_anthropic_protected(&input.bytes, policy),
367        InputFormat::GitDiff => extract_diff_protected(&input.bytes),
368        // ponytail: no transform touches plain text/command output structure yet beyond
369        // log/diff compaction (task-scope gated), so nothing is unconditionally protected.
370        // Generic Json has no "protected" sub-segment either — json_field_fold's own
371        // round-trip safety gate is what guarantees its data is preserved.
372        InputFormat::PlainText
373        | InputFormat::CommandOutput
374        | InputFormat::Json
375        | InputFormat::Auto => Vec::new(),
376    }
377}
378
379fn extract_openai_protected(bytes: &[u8], policy: &CompressionPolicy) -> Vec<Vec<u8>> {
380    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
381        return Vec::new();
382    };
383    let Some(messages) = value.get("messages").and_then(|m| m.as_array()) else {
384        return Vec::new();
385    };
386
387    let mut segments = Vec::new();
388    for message in messages {
389        if message.get("role").and_then(|r| r.as_str()) == Some("system")
390            && let Some(bytes) = message_content_bytes(message)
391        {
392            segments.push(bytes);
393        }
394    }
395    if policy.preserve_latest_user_message
396        && let Some(last_user) = messages
397            .iter()
398            .rev()
399            .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
400        && let Some(bytes) = message_content_bytes(last_user)
401    {
402        segments.push(bytes);
403    }
404    segments
405}
406
407fn extract_anthropic_protected(bytes: &[u8], policy: &CompressionPolicy) -> Vec<Vec<u8>> {
408    let Ok(value) = serde_json::from_slice::<serde_json::Value>(bytes) else {
409        return Vec::new();
410    };
411
412    let mut segments = Vec::new();
413    // Anthropic's `system` field is either a plain string OR a structured array of content
414    // blocks (`[{"type":"text","text":"..."}, ...]`) -- the structured shape used to fall
415    // through `.as_str()` as `None` and get zero protection. Mirrors `message_content_bytes`'s
416    // existing string-or-structured handling for message `content`.
417    match value.get("system") {
418        Some(serde_json::Value::String(text)) => segments.push(text.as_bytes().to_vec()),
419        Some(structured @ serde_json::Value::Array(_)) => {
420            if let Ok(bytes) = serde_json::to_vec(structured) {
421                segments.push(bytes);
422            }
423        }
424        _ => {}
425    }
426    if policy.preserve_latest_user_message
427        && let Some(last_user) =
428            value
429                .get("messages")
430                .and_then(|m| m.as_array())
431                .and_then(|messages| {
432                    messages
433                        .iter()
434                        .rev()
435                        .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
436                })
437        && let Some(bytes) = message_content_bytes(last_user)
438    {
439        segments.push(bytes);
440    }
441    segments
442}
443
444fn message_content_bytes(message: &serde_json::Value) -> Option<Vec<u8>> {
445    match message.get("content") {
446        Some(serde_json::Value::String(text)) => Some(text.as_bytes().to_vec()),
447        Some(structured) => serde_json::to_vec(structured).ok(),
448        None => None,
449    }
450}
451
452/// Keeps file names and hunk headers, matching the `diff_compaction` contract of what must
453/// survive compaction. Each kept line is its own segment.
454fn extract_diff_protected(bytes: &[u8]) -> Vec<Vec<u8>> {
455    let text = String::from_utf8_lossy(bytes);
456    let mut segments = Vec::new();
457    for line in text.lines() {
458        if line.starts_with("diff --git")
459            || line.starts_with("--- ")
460            || line.starts_with("+++ ")
461            || line.starts_with("@@")
462        {
463            let mut segment = line.as_bytes().to_vec();
464            segment.push(b'\n');
465            segments.push(segment);
466        }
467    }
468    segments
469}
470
471#[cfg(test)]
472mod tests {
473    use super::*;
474    use crate::token_estimator::ByteHeuristicEstimator;
475
476    #[test]
477    fn default_mode_is_balanced() {
478        let policy = CompressionPolicy::builder().build().unwrap();
479        assert_eq!(policy.preset, Preset::Balanced);
480    }
481
482    #[test]
483    fn store_originals_defaults_to_false_with_a_default_namespace() {
484        let policy = CompressionPolicy::builder().build().unwrap();
485        assert!(!policy.store_originals);
486        assert_eq!(policy.retrieval_namespace, "default");
487    }
488
489    #[test]
490    fn store_originals_and_namespace_are_settable_via_the_builder() {
491        let policy = CompressionPolicy::builder()
492            .store_originals(true)
493            .retrieval_namespace("project-x")
494            .retrieval_ttl_seconds(Some(60))
495            .retrieval_backend("memory")
496            .retrieval_store_path(Some(std::path::PathBuf::from("/tmp/custom")))
497            .build()
498            .unwrap();
499        assert!(policy.store_originals);
500        assert_eq!(policy.retrieval_namespace, "project-x");
501        assert_eq!(policy.retrieval_ttl_seconds, Some(60));
502        assert_eq!(policy.retrieval_backend, "memory");
503        assert_eq!(
504            policy.retrieval_store_path,
505            Some(std::path::PathBuf::from("/tmp/custom"))
506        );
507    }
508
509    #[test]
510    fn retrieval_defaults_are_none_ttl_and_filesystem_backend() {
511        let policy = CompressionPolicy::builder().build().unwrap();
512        assert_eq!(policy.retrieval_ttl_seconds, None);
513        assert_eq!(policy.retrieval_backend, "filesystem");
514        assert_eq!(policy.retrieval_store_path, None);
515    }
516
517    #[test]
518    fn secret_redaction_cannot_be_disabled_through_policy() {
519        let err = CompressionPolicy::builder()
520            .disable("secret_redaction")
521            .build()
522            .unwrap_err();
523        assert!(matches!(err, TokenFoldError::ConfigError(_)));
524    }
525
526    #[test]
527    fn disabling_other_transforms_is_allowed() {
528        let policy = CompressionPolicy::builder()
529            .disable("json_minify")
530            .build()
531            .unwrap();
532        assert_eq!(policy.disabled, vec!["json_minify".to_string()]);
533    }
534
535    #[test]
536    fn floor_is_zero_for_plain_text() {
537        let input = CompressionInput::plain_text(b"just some plain text".to_vec());
538        let policy = CompressionPolicy::builder().build().unwrap();
539        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
540        assert_eq!(floor, 0);
541    }
542
543    #[test]
544    fn floor_covers_system_and_latest_user_message_for_openai_json() {
545        let payload = serde_json::json!({
546            "model": "gpt-4",
547            "messages": [
548                {"role": "system", "content": "You are a helpful assistant."},
549                {"role": "user", "content": "first question"},
550                {"role": "assistant", "content": "first answer"},
551                {"role": "user", "content": "second question"},
552            ]
553        });
554        let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
555        let policy = CompressionPolicy::builder().build().unwrap();
556        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
557
558        let expected_bytes = "You are a helpful assistant.".len() + "second question".len();
559        assert_eq!(
560            floor,
561            ByteHeuristicEstimator.count_bytes(&vec![0u8; expected_bytes])
562        );
563        // The earlier "first question" turn must NOT be counted as protected.
564        assert!(floor < ByteHeuristicEstimator.count_bytes(input.bytes.as_slice()));
565    }
566
567    #[test]
568    fn floor_excludes_latest_user_message_when_policy_disables_preservation() {
569        let payload = serde_json::json!({
570            "messages": [
571                {"role": "system", "content": "system prompt"},
572                {"role": "user", "content": "question"},
573            ]
574        });
575        let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
576        let policy = CompressionPolicy::builder()
577            .preserve_latest_user_message(false)
578            .build()
579            .unwrap();
580        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
581        assert_eq!(floor, ByteHeuristicEstimator.count_bytes(b"system prompt"));
582    }
583
584    #[test]
585    fn floor_covers_system_and_latest_user_message_for_anthropic_json() {
586        let payload = serde_json::json!({
587            "system": "system prompt",
588            "messages": [
589                {"role": "user", "content": "first"},
590                {"role": "assistant", "content": "reply"},
591                {"role": "user", "content": "second"},
592            ]
593        });
594        let input = CompressionInput::anthropic_json(serde_json::to_vec(&payload).unwrap());
595        let policy = CompressionPolicy::builder().build().unwrap();
596        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
597        let expected_bytes = "system prompt".len() + "second".len();
598        assert_eq!(
599            floor,
600            ByteHeuristicEstimator.count_bytes(&vec![0u8; expected_bytes])
601        );
602    }
603
604    #[test]
605    fn floor_covers_structured_anthropic_system_content_not_just_a_plain_string() {
606        // Round-4 external review: Anthropic's `system` field can be a structured array of
607        // content blocks (`[{"type":"text","text":"..."}]`), not just a plain string --
608        // `.as_str()` alone returned `None` for that shape, so a structured system prompt got
609        // ZERO protection (silently prunable/rewritable like any other content).
610        let payload = serde_json::json!({
611            "system": [{"type": "text", "text": "structured system prompt"}],
612            "messages": [
613                {"role": "user", "content": "first"},
614            ]
615        });
616        let input = CompressionInput::anthropic_json(serde_json::to_vec(&payload).unwrap());
617        let policy = CompressionPolicy::builder()
618            .preserve_latest_user_message(false)
619            .build()
620            .unwrap();
621        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
622        assert!(
623            floor > 0,
624            "a structured Anthropic `system` array must contribute to the protected floor"
625        );
626        let segments = protected_segments(&input, &policy);
627        let system_bytes = serde_json::to_vec(
628            &serde_json::json!([{"type": "text", "text": "structured system prompt"}]),
629        )
630        .unwrap();
631        assert!(
632            segments.contains(&system_bytes),
633            "the structured system content must be a protected segment, byte-for-byte"
634        );
635    }
636
637    #[test]
638    fn floor_keeps_diff_headers_and_hunk_markers_only() {
639        let diff =
640            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";
641        let input = CompressionInput::git_diff(diff.to_vec());
642        let policy = CompressionPolicy::builder().build().unwrap();
643        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
644        assert!(floor > 0);
645        assert!(floor < ByteHeuristicEstimator.count_bytes(diff));
646    }
647
648    #[test]
649    fn lossy_defaults_to_disabled_with_a_default_ratio() {
650        let policy = CompressionPolicy::builder().build().unwrap();
651        assert_eq!(policy.lossy, None);
652        assert_eq!(policy.lossy_ratio, 0.3);
653        assert!(policy.lossy_preserve.is_empty());
654    }
655
656    #[test]
657    fn lossy_is_settable_via_the_builder() {
658        let policy = CompressionPolicy::builder()
659            .lossy(LossyPath::Heuristic)
660            .lossy_ratio(0.5)
661            .lossy_preserve("items")
662            .lossy_preserve("data.results")
663            .build()
664            .unwrap();
665        assert_eq!(policy.lossy, Some(LossyPath::Heuristic));
666        assert_eq!(policy.lossy_ratio, 0.5);
667        assert_eq!(policy.lossy_preserve, vec!["items", "data.results"]);
668    }
669
670    #[test]
671    fn lossy_refuses_memory_retrieval_backend() {
672        let err = CompressionPolicy::builder()
673            .lossy(LossyPath::Heuristic)
674            .retrieval_backend("memory")
675            .build()
676            .unwrap_err();
677        assert!(matches!(err, TokenFoldError::ConfigError(_)));
678    }
679
680    #[test]
681    fn lossy_refuses_a_ttl_below_the_floor() {
682        let err = CompressionPolicy::builder()
683            .lossy(LossyPath::Heuristic)
684            .retrieval_ttl_seconds(Some(60))
685            .build()
686            .unwrap_err();
687        assert!(matches!(err, TokenFoldError::ConfigError(_)));
688    }
689
690    #[test]
691    fn lossy_with_default_retrieval_settings_is_accepted() {
692        // Defaults (filesystem backend, 7-day TTL) already clear the floor -- a user shouldn't
693        // need to configure retrieval explicitly just to use --lossy.
694        let policy = CompressionPolicy::builder()
695            .lossy(LossyPath::Heuristic)
696            .build()
697            .unwrap();
698        assert_eq!(policy.lossy, Some(LossyPath::Heuristic));
699    }
700
701    #[test]
702    fn lossy_ratio_outside_unit_interval_is_rejected() {
703        let err = CompressionPolicy::builder()
704            .lossy_ratio(1.5)
705            .build()
706            .unwrap_err();
707        assert!(matches!(err, TokenFoldError::ConfigError(_)));
708        let err = CompressionPolicy::builder()
709            .lossy_ratio(-0.1)
710            .build()
711            .unwrap_err();
712        assert!(matches!(err, TokenFoldError::ConfigError(_)));
713    }
714
715    #[test]
716    fn malformed_json_never_panics_and_yields_zero_floor() {
717        let input = CompressionInput::openai_json(b"{not json".to_vec());
718        let policy = CompressionPolicy::builder().build().unwrap();
719        let floor = protected_floor(&input, &policy, &ByteHeuristicEstimator);
720        assert_eq!(floor, 0);
721    }
722}
tokenfold_core/budget.rs

tokenfold_core/
budget.rs
