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

tokenfold_core/
budget.rs
