1use std::collections::HashMap;
2
3use serde_json::Value;
4
5use crate::budget::{CompressionMode, CompressionPolicy, TaskScope, protected_segments};
6use crate::errors::TokenFoldError;
7use crate::input::{CompressionInput, CompressionOutput, InputFormat};
8use crate::modes::{self, ModeEntry, TransformId};
9use crate::report::{
10 BudgetReport, CompressionReport, QualityReport, RetrievalReport, Severity, SkippedReason,
11 TransformReport, TransformStatus, Warning, WarningCode,
12};
13use crate::retrieval_store::{self, RetrievalStore};
14use crate::safety;
15use crate::status::Status;
16use crate::token_estimator::{ByteHeuristicEstimator, TokenEstimator};
17use crate::transforms;
18
19pub fn compress(
22 input: CompressionInput,
23 policy: &CompressionPolicy,
24) -> Result<CompressionOutput, TokenFoldError> {
25 #[cfg(feature = "tiktoken")]
26 {
27 if let Ok(estimator) = crate::token_estimator::TiktokenEstimator::o200k_base() {
28 return compress_with_estimator(input, policy, &estimator);
29 }
30 }
31 compress_with_estimator(input, policy, &ByteHeuristicEstimator)
32}
33
34pub fn compress_with_estimator(
35 input: CompressionInput,
36 policy: &CompressionPolicy,
37 estimator: &dyn TokenEstimator,
38) -> Result<CompressionOutput, TokenFoldError> {
39 policy.validate()?;
44 let original_tokens = estimator.count_bytes(&input.bytes);
45 let target = policy.target_tokens;
46 let estimator_info = estimator.info();
47
48 let retrieval = maybe_store_originals(&input.bytes, input.format, policy);
51
52 if let Some(t) = target
55 && original_tokens <= t
56 {
57 let mut warnings = Vec::new();
58 if !estimator_info.is_exact {
59 warnings.push(heuristic_budget_warning());
60 }
61 let mut report = CompressionReport::new(
62 original_tokens,
63 original_tokens,
64 estimator_info,
65 Status::Passthrough,
66 mode_label(policy.mode).to_string(),
67 format_label(input.format).to_string(),
68 task_scope_label(policy.task_scope).to_string(),
69 Vec::new(),
70 warnings,
71 );
72 report.retrieval = retrieval;
73 return Ok(CompressionOutput {
74 bytes: input.bytes,
75 report,
76 });
77 }
78
79 apply_transforms(input, policy, estimator, original_tokens, target, retrieval)
80}
81
82fn maybe_store_originals(
89 input_bytes: &[u8],
90 format: InputFormat,
91 policy: &CompressionPolicy,
92) -> Option<RetrievalReport> {
93 if policy.preview {
98 return None;
99 }
100 let lossy_can_run = policy.lossy.is_some() && format == InputFormat::Json;
111 if !policy.store_originals && !lossy_can_run {
112 return None;
113 }
114 let ttl_seconds = policy
115 .retrieval_ttl_seconds
116 .unwrap_or(retrieval_store::DEFAULT_TTL_SECONDS);
117 let skipped = || RetrievalReport {
118 store_namespace: policy.retrieval_namespace.clone(),
119 hash_algorithm: "sha256".to_string(),
120 marker_count: 0,
121 ttl_seconds: None,
122 persisted_original_bytes: 0,
123 skipped_original_bytes: input_bytes.len(),
124 };
125 let Ok(store) = RetrievalStore::open(
126 &policy.retrieval_backend,
127 "sha256",
128 policy.retrieval_store_path.clone(),
129 ) else {
130 return Some(skipped());
131 };
132 Some(
133 match store.store(input_bytes, &policy.retrieval_namespace, Some(ttl_seconds)) {
134 Ok(_marker) => RetrievalReport {
135 store_namespace: policy.retrieval_namespace.clone(),
136 hash_algorithm: "sha256".to_string(),
137 marker_count: 1,
138 ttl_seconds: Some(ttl_seconds),
139 persisted_original_bytes: input_bytes.len(),
140 skipped_original_bytes: 0,
141 },
142 Err(_) => skipped(),
143 },
144 )
145}
146
147fn apply_transforms(
148 input: CompressionInput,
149 policy: &CompressionPolicy,
150 estimator: &dyn TokenEstimator,
151 original_tokens: usize,
152 target: Option<usize>,
153 retrieval: Option<RetrievalReport>,
154) -> Result<CompressionOutput, TokenFoldError> {
155 let estimator_info = estimator.info();
156 let mut warnings = Vec::new();
157 let mut transform_reports = Vec::new();
158 if !estimator_info.is_exact {
159 warnings.push(heuristic_budget_warning());
160 }
161
162 let mut bytes;
166 if policy.unsafe_disable_redaction {
167 bytes = input.bytes.clone();
168 warnings.push(Warning {
169 code: WarningCode::UnredactedContentPossible,
170 severity: Severity::Critical,
171 transform: Some("secret_redaction".to_string()),
172 message: "redaction was disabled via unsafe_disable_redaction; output may contain unredacted secrets".to_string(),
173 });
174 transform_reports.push(skipped_at(
175 "secret_redaction",
176 "1.0.0",
177 original_tokens,
178 SkippedReason::DisabledByUser,
179 ));
180 } else {
181 let outcome = transforms::redaction::redact(&input.bytes);
182 let tokens_after = estimator.count_bytes(&outcome.bytes);
183 warnings.push(Warning {
184 code: WarningCode::UnredactedContentPossible,
185 severity: Severity::Info,
186 transform: Some("secret_redaction".to_string()),
187 message: "redaction is best-effort; it is not a guarantee that no secret survives"
188 .to_string(),
189 });
190 transform_reports.push(TransformReport {
191 id: "secret_redaction".to_string(),
192 version: "1.0.0".to_string(),
193 tokens_before: original_tokens,
194 tokens_after,
195 saved_tokens: original_tokens.saturating_sub(tokens_after),
196 savings_ratio: ratio(original_tokens, tokens_after),
197 elapsed_micros: None,
198 status: if outcome.redacted_count > 0 {
199 TransformStatus::Applied
200 } else {
201 TransformStatus::NoOp
202 },
203 skipped_reason: None,
204 warnings: Vec::new(),
205 });
206 bytes = outcome.bytes;
207 }
208
209 let working_input = CompressionInput {
213 format: input.format,
214 bytes: bytes.clone(),
215 };
216 let protected = protected_segments(&working_input, policy);
217 let floor = estimator.count_bytes(&protected.concat());
218 let mut current_tokens = estimator.count_bytes(&bytes);
219
220 if let Some(t) = target
221 && t < floor
222 {
223 warnings.push(Warning {
224 code: WarningCode::UnreachableTarget,
225 severity: Severity::Warn,
226 transform: None,
227 message: format!("target {t} tokens is below the protected floor of {floor} tokens"),
228 });
229 let mut report = CompressionReport::new(
230 original_tokens,
231 current_tokens,
232 estimator_info,
233 Status::UnreachableTarget,
234 mode_label(policy.mode).to_string(),
235 format_label(input.format).to_string(),
236 task_scope_label(policy.task_scope).to_string(),
237 transform_reports,
238 warnings,
239 );
240 report.budget = Some(BudgetReport {
241 target_tokens: target,
242 protected_floor: floor,
243 achieved_tokens: current_tokens,
244 });
245 report.retrieval = retrieval;
246 return Ok(CompressionOutput { bytes, report });
247 }
248
249 let defer_for_lossy = |entry: &ModeEntry| {
268 policy.lossy.is_some()
269 && matches!(
270 entry.transform_id,
271 TransformId::JsonFieldFold | TransformId::JsonValueDict
272 )
273 };
274 let entries = modes::pipeline_for(
275 policy.mode,
276 policy.task_scope,
277 input.format,
278 policy.experimental,
279 &policy.enable,
280 &policy.disabled,
281 );
282 let mut deferred: Vec<&ModeEntry> = Vec::new();
283 for entry in entries {
284 if defer_for_lossy(entry) {
285 deferred.push(entry);
286 continue;
287 }
288 run_transform_entry(
289 entry,
290 policy,
291 input.format,
292 estimator,
293 &protected,
294 target,
295 &mut bytes,
296 &mut current_tokens,
297 &mut transform_reports,
298 &mut warnings,
299 );
300 }
301
302 let mut retrieval = retrieval;
316 let mut lossy_applied = false;
317 if policy.lossy.is_some() {
318 let mut lossless_bytes = bytes.clone();
321 let mut lossless_tokens = current_tokens;
322 let mut lossless_reports = Vec::new();
323 let mut lossless_warnings = Vec::new();
324 for entry in &deferred {
325 run_transform_entry(
326 entry,
327 policy,
328 input.format,
329 estimator,
330 &protected,
331 target,
332 &mut lossless_bytes,
333 &mut lossless_tokens,
334 &mut lossless_reports,
335 &mut lossless_warnings,
336 );
337 }
338
339 let target_met_losslessly = target.is_some_and(|t| lossless_tokens <= t);
348 let lossy_report = if target_met_losslessly {
349 skipped_at(
350 transforms::json_prune::TRANSFORM_ID,
351 transforms::json_prune::TRANSFORM_VERSION,
352 current_tokens,
353 SkippedReason::TargetAlreadyMet,
354 )
355 } else {
356 let (new_bytes, new_tokens, lossy_report) = apply_lossy_reduction(
357 &bytes,
358 current_tokens,
359 lossless_tokens,
360 policy,
361 input.format,
362 estimator,
363 &protected,
364 &mut retrieval,
365 )?;
366 lossy_applied = lossy_report.status == TransformStatus::Applied;
367 if lossy_applied {
368 bytes = new_bytes;
369 current_tokens = new_tokens;
370 }
371 lossy_report
372 };
373 transform_reports.push(lossy_report);
374
375 if lossy_applied {
376 for entry in deferred {
381 transform_reports.push(skipped(
382 entry,
383 current_tokens,
384 SkippedReason::IncompatibleWithLossy,
385 ));
386 }
387 } else {
388 bytes = lossless_bytes;
392 current_tokens = lossless_tokens;
393 transform_reports.extend(lossless_reports);
394 warnings.extend(lossless_warnings);
395 }
396 }
397
398 let status = match target {
399 None => Status::BestEffort,
400 Some(t) if current_tokens <= t => Status::Compressed,
401 Some(_) => Status::BestEffort,
402 };
403
404 let mut report = CompressionReport::new(
405 original_tokens,
406 current_tokens,
407 estimator_info,
408 status,
409 mode_label(policy.mode).to_string(),
410 format_label(input.format).to_string(),
411 task_scope_label(policy.task_scope).to_string(),
412 transform_reports,
413 warnings,
414 );
415 report.budget = Some(BudgetReport {
416 target_tokens: target,
417 protected_floor: floor,
418 achieved_tokens: current_tokens,
419 });
420 if lossy_applied {
427 report.quality = Some(QualityReport {
428 eval_profile_id: "unvalidated".to_string(),
429 task_scope: task_scope_label(policy.task_scope).to_string(),
430 validated_ratio_band: None,
431 quality_retention: None,
432 contrastive_failure_rate: None,
433 gate_passed: false,
434 });
435 }
436 report.retrieval = retrieval;
437 Ok(CompressionOutput { bytes, report })
438}
439
440#[allow(clippy::too_many_arguments)]
445fn run_transform_entry(
446 entry: &ModeEntry,
447 policy: &CompressionPolicy,
448 format: InputFormat,
449 estimator: &dyn TokenEstimator,
450 protected: &[Vec<u8>],
451 target: Option<usize>,
452 bytes: &mut Vec<u8>,
453 current_tokens: &mut usize,
454 transform_reports: &mut Vec<TransformReport>,
455 warnings: &mut Vec<Warning>,
456) {
457 if let Some(t) = target
458 && *current_tokens <= t
459 {
460 transform_reports.push(skipped(
461 entry,
462 *current_tokens,
463 SkippedReason::TargetAlreadyMet,
464 ));
465 return;
466 }
467
468 let tokens_before = *current_tokens;
469 let max_ratio = entry.max_ratio_for(policy.mode);
470
471 let candidate = match apply_single_transform(entry.transform_id, bytes, policy) {
472 Ok(candidate) => candidate,
473 Err(_) => {
474 transform_reports.push(skipped(
475 entry,
476 tokens_before,
477 SkippedReason::NotApplicableToFormat,
478 ));
479 return;
480 }
481 };
482
483 let tokens_after_candidate = estimator.count_bytes(&candidate);
484 if tokens_after_candidate > tokens_before {
485 transform_reports.push(skipped(
487 entry,
488 tokens_before,
489 SkippedReason::WouldIncreaseTokens,
490 ));
491 return;
492 }
493 if tokens_after_candidate == tokens_before {
494 transform_reports.push(TransformReport {
497 id: entry.transform_id.as_str().to_string(),
498 version: entry.version.to_string(),
499 tokens_before,
500 tokens_after: tokens_before,
501 saved_tokens: 0,
502 savings_ratio: 0.0,
503 elapsed_micros: None,
504 status: TransformStatus::NoOp,
505 skipped_reason: None,
506 warnings: Vec::new(),
507 });
508 return;
509 }
510 let ratio_used = 1.0 - (tokens_after_candidate as f64 / tokens_before.max(1) as f64);
511 if ratio_used > max_ratio {
512 transform_reports.push(skipped(
513 entry,
514 tokens_before,
515 SkippedReason::NotEnabledInMode,
516 ));
517 return;
518 }
519
520 if !validate_safety(entry.transform_id, format, bytes, &candidate, protected) {
521 transform_reports.push(rolled_back(entry, tokens_before));
522 warnings.push(safety_downgrade_warning(entry.transform_id.as_str()));
523 return;
524 }
525
526 *bytes = candidate;
527 *current_tokens = tokens_after_candidate;
528 transform_reports.push(TransformReport {
529 id: entry.transform_id.as_str().to_string(),
530 version: entry.version.to_string(),
531 tokens_before,
532 tokens_after: *current_tokens,
533 saved_tokens: tokens_before.saturating_sub(*current_tokens),
534 savings_ratio: ratio(tokens_before, *current_tokens),
535 elapsed_micros: None,
536 status: TransformStatus::Applied,
537 skipped_reason: None,
538 warnings: Vec::new(),
539 });
540}
541
542fn apply_single_transform(
543 transform_id: TransformId,
544 bytes: &[u8],
545 policy: &CompressionPolicy,
546) -> Result<Vec<u8>, String> {
547 match transform_id {
548 TransformId::JsonMinify => transforms::json::minify_json(bytes).map_err(|e| e.to_string()),
549 TransformId::JsonFieldFold => {
550 transforms::json_fold::fold_json(bytes).map_err(|e| e.to_string())
551 }
552 TransformId::JsonValueDict => {
553 transforms::json_dict::dict_json(bytes).map_err(|e| e.to_string())
554 }
555 TransformId::SchemaCompaction => {
556 transforms::schema::compact_schema(bytes, 1).map_err(|e| e.to_string())
560 }
561 TransformId::LogFieldFold => {
562 let text = std::str::from_utf8(bytes).map_err(|e| e.to_string())?;
563 Ok(transforms::log_fold::fold_log(text).into_bytes())
564 }
565 TransformId::LogCompaction => {
566 let text = std::str::from_utf8(bytes).map_err(|e| e.to_string())?;
567 Ok(transforms::logs::compact(text, false).into_bytes())
568 }
569 TransformId::DiffCompaction => {
570 let text = std::str::from_utf8(bytes).map_err(|e| e.to_string())?;
571 let keep_line_bodies = policy.task_scope != TaskScope::ChangeSummary;
572 Ok(transforms::diff::compact_diff(text, keep_line_bodies).into_bytes())
573 }
574 }
575}
576
577#[allow(clippy::too_many_arguments)]
583fn apply_lossy_reduction(
584 bytes: &[u8],
585 tokens_before: usize,
586 must_beat_tokens: usize,
587 policy: &CompressionPolicy,
588 format: InputFormat,
589 estimator: &dyn TokenEstimator,
590 protected: &[Vec<u8>],
591 retrieval: &mut Option<RetrievalReport>,
592) -> Result<(Vec<u8>, usize, TransformReport), TokenFoldError> {
593 let noop = |status, reason| TransformReport {
594 id: transforms::json_prune::TRANSFORM_ID.to_string(),
595 version: transforms::json_prune::TRANSFORM_VERSION.to_string(),
596 tokens_before,
597 tokens_after: tokens_before,
598 saved_tokens: 0,
599 savings_ratio: 0.0,
600 elapsed_micros: None,
601 status,
602 skipped_reason: reason,
603 warnings: Vec::new(),
604 };
605
606 if format != InputFormat::Json {
617 return Ok((
618 bytes.to_vec(),
619 tokens_before,
620 noop(
621 TransformStatus::Skipped,
622 Some(SkippedReason::NotApplicableToFormat),
623 ),
624 ));
625 }
626
627 let options = transforms::json_prune::LossyOptions {
628 preserve_paths: policy.lossy_preserve.clone(),
629 ratio: policy.lossy_ratio,
630 namespace: policy.retrieval_namespace.clone(),
631 };
632 let outcome = match transforms::json_prune::prune(bytes, &options, estimator) {
633 Ok(Some(outcome)) => outcome,
634 Ok(None) => {
635 return Ok((
636 bytes.to_vec(),
637 tokens_before,
638 noop(TransformStatus::NoOp, None),
639 ));
640 }
641 Err(_) => {
644 return Ok((
645 bytes.to_vec(),
646 tokens_before,
647 noop(
648 TransformStatus::Skipped,
649 Some(SkippedReason::NotApplicableToFormat),
650 ),
651 ));
652 }
653 };
654
655 let best_case_bytes = serde_json::to_vec(&outcome.json).map_err(|e| {
669 TokenFoldError::InternalError(format!("failed to serialize pruned json: {e}"))
670 })?;
671 if estimator.count_bytes(&best_case_bytes) >= must_beat_tokens {
672 return Ok((
673 bytes.to_vec(),
674 tokens_before,
675 noop(TransformStatus::RolledBack, None),
676 ));
677 }
678
679 let ttl_seconds = policy
680 .retrieval_ttl_seconds
681 .unwrap_or(retrieval_store::DEFAULT_TTL_SECONDS);
682 let store = if policy.preview {
690 None
691 } else {
692 RetrievalStore::open(
693 &policy.retrieval_backend,
694 "sha256",
695 policy.retrieval_store_path.clone(),
696 )
697 .ok()
698 };
699
700 let preview_probe = policy.preview.then(RetrievalStore::memory);
707
708 let mut restore: HashMap<String, Value> = HashMap::new();
709 let mut persisted_bytes = 0usize;
710 let mut skipped_bytes = 0usize;
711 let mut marker_count = 0usize;
712 for item in &outcome.dropped {
713 if let Some(probe) = &preview_probe {
714 let would_store = probe
715 .store(&item.bytes, &policy.retrieval_namespace, Some(ttl_seconds))
716 .is_ok();
717 if !would_store {
718 let original: Value = serde_json::from_slice(&item.bytes).map_err(|e| {
721 TokenFoldError::InternalError(format!(
722 "json_prune produced a dropped item that isn't valid JSON: {e}"
723 ))
724 })?;
725 restore.insert(item.hash.clone(), original);
726 }
727 continue;
731 }
732 let stored = store.as_ref().and_then(|s| {
733 s.store(&item.bytes, &policy.retrieval_namespace, Some(ttl_seconds))
734 .ok()
735 });
736 match stored {
737 Some(_marker) => {
738 persisted_bytes += item.bytes.len();
739 marker_count += 1;
740 }
741 None => {
742 let original: Value = serde_json::from_slice(&item.bytes).map_err(|e| {
743 TokenFoldError::InternalError(format!(
744 "json_prune produced a dropped item that isn't valid JSON: {e}"
745 ))
746 })?;
747 restore.insert(item.hash.clone(), original);
748 skipped_bytes += item.bytes.len();
749 }
750 }
751 }
752
753 let final_json = transforms::json_prune::revert_markers(&outcome.json, &restore);
754 let final_bytes = serde_json::to_vec(&final_json).map_err(|e| {
755 TokenFoldError::InternalError(format!("failed to serialize pruned json: {e}"))
756 })?;
757 let tokens_after_candidate = estimator.count_bytes(&final_bytes);
758
759 let violates_protected_segments = !safety::protected_segments_present(protected, &final_bytes);
776 let status = if final_bytes == bytes {
781 TransformStatus::NoOp
782 } else if tokens_after_candidate >= must_beat_tokens || violates_protected_segments {
783 TransformStatus::RolledBack
784 } else {
785 TransformStatus::Applied
786 };
787 let bytes_out = if status == TransformStatus::RolledBack {
788 bytes.to_vec()
789 } else {
790 final_bytes
791 };
792 let tokens_after = if status == TransformStatus::RolledBack {
793 tokens_before
794 } else {
795 tokens_after_candidate
796 };
797
798 if status != TransformStatus::RolledBack {
799 match retrieval {
800 Some(existing) => {
801 existing.marker_count += marker_count;
802 existing.persisted_original_bytes += persisted_bytes;
803 existing.skipped_original_bytes += skipped_bytes;
804 if marker_count > 0 {
813 existing.ttl_seconds = Some(ttl_seconds);
814 }
815 }
816 None if marker_count > 0 || skipped_bytes > 0 => {
817 *retrieval = Some(RetrievalReport {
818 store_namespace: policy.retrieval_namespace.clone(),
819 hash_algorithm: "sha256".to_string(),
820 marker_count,
821 ttl_seconds: Some(ttl_seconds),
822 persisted_original_bytes: persisted_bytes,
823 skipped_original_bytes: skipped_bytes,
824 });
825 }
826 None => {}
827 }
828 }
829
830 let report_warnings = if status == TransformStatus::RolledBack && violates_protected_segments {
831 vec![safety_downgrade_warning(
832 transforms::json_prune::TRANSFORM_ID,
833 )]
834 } else {
835 Vec::new()
836 };
837
838 Ok((
839 bytes_out,
840 tokens_after,
841 TransformReport {
842 id: transforms::json_prune::TRANSFORM_ID.to_string(),
843 version: transforms::json_prune::TRANSFORM_VERSION.to_string(),
844 tokens_before,
845 tokens_after,
846 saved_tokens: tokens_before.saturating_sub(tokens_after),
847 savings_ratio: ratio(tokens_before, tokens_after),
848 elapsed_micros: None,
849 status,
850 skipped_reason: None,
851 warnings: report_warnings,
852 },
853 ))
854}
855
856fn validate_safety(
857 transform_id: TransformId,
858 format: InputFormat,
859 before: &[u8],
860 after: &[u8],
861 protected: &[Vec<u8>],
862) -> bool {
863 match transform_id {
864 TransformId::JsonFieldFold => {
868 if !safety::json_still_valid(after) {
869 return false;
870 }
871 if !transforms::json_fold::round_trips(before, after) {
872 return false;
873 }
874 }
875 TransformId::JsonValueDict => {
878 if !safety::json_still_valid(after) {
879 return false;
880 }
881 if !transforms::json_dict::round_trips(before, after) {
882 return false;
883 }
884 }
885 TransformId::JsonMinify | TransformId::SchemaCompaction => {
888 let is_json_format = matches!(
889 format,
890 InputFormat::OpenAiJson | InputFormat::AnthropicJson | InputFormat::Json
891 );
892 if is_json_format {
893 if !safety::json_still_valid(after) {
894 return false;
895 }
896 if !safety::json_key_order_preserved(before, after) {
897 return false;
898 }
899 }
900 }
901 TransformId::LogFieldFold => {
904 if !transforms::log_fold::round_trips(before, after) {
905 return false;
906 }
907 }
908 TransformId::LogCompaction | TransformId::DiffCompaction => {}
909 }
910 safety::protected_segments_present(protected, after)
911}
912
913fn skipped(entry: &ModeEntry, tokens: usize, reason: SkippedReason) -> TransformReport {
914 skipped_at(entry.transform_id.as_str(), entry.version, tokens, reason)
915}
916
917fn skipped_at(id: &str, version: &str, tokens: usize, reason: SkippedReason) -> TransformReport {
918 TransformReport {
919 id: id.to_string(),
920 version: version.to_string(),
921 tokens_before: tokens,
922 tokens_after: tokens,
923 saved_tokens: 0,
924 savings_ratio: 0.0,
925 elapsed_micros: None,
926 status: TransformStatus::Skipped,
927 skipped_reason: Some(reason),
928 warnings: Vec::new(),
929 }
930}
931
932fn rolled_back(entry: &ModeEntry, tokens: usize) -> TransformReport {
933 TransformReport {
934 id: entry.transform_id.as_str().to_string(),
935 version: entry.version.to_string(),
936 tokens_before: tokens,
937 tokens_after: tokens,
938 saved_tokens: 0,
939 savings_ratio: 0.0,
940 elapsed_micros: None,
941 status: TransformStatus::RolledBack,
942 skipped_reason: None,
943 warnings: Vec::new(),
944 }
945}
946
947fn safety_downgrade_warning(transform_id: &str) -> Warning {
948 Warning {
949 code: WarningCode::SafetyDowngrade,
950 severity: Severity::Warn,
951 transform: Some(transform_id.to_string()),
952 message: format!(
953 "{transform_id} was rolled back: a safety invariant would have been violated"
954 ),
955 }
956}
957
958fn heuristic_budget_warning() -> Warning {
959 Warning {
960 code: WarningCode::HeuristicBudgetUsed,
961 severity: Severity::Info,
962 transform: None,
963 message: "token counts are heuristic estimates (~bytes/4), not exact".to_string(),
964 }
965}
966
967fn ratio(before: usize, after: usize) -> f64 {
968 if before == 0 {
969 0.0
970 } else {
971 before.saturating_sub(after) as f64 / before as f64
972 }
973}
974
975fn mode_label(mode: CompressionMode) -> &'static str {
976 match mode {
977 CompressionMode::Conservative => "conservative",
978 CompressionMode::Balanced => "balanced",
979 CompressionMode::Aggressive => "aggressive",
980 }
981}
982
983fn format_label(format: InputFormat) -> &'static str {
984 match format {
985 InputFormat::Auto => "auto",
986 InputFormat::OpenAiJson => "openai_json",
987 InputFormat::AnthropicJson => "anthropic_json",
988 InputFormat::Json => "json",
989 InputFormat::PlainText => "plain_text",
990 InputFormat::CommandOutput => "command_output",
991 InputFormat::GitDiff => "git_diff",
992 }
993}
994
995fn task_scope_label(scope: TaskScope) -> &'static str {
996 match scope {
997 TaskScope::All => "all",
998 TaskScope::General => "general",
999 TaskScope::CodeReview => "code_review",
1000 TaskScope::ChangeSummary => "change_summary",
1001 TaskScope::Debugging => "debugging",
1002 TaskScope::Generation => "generation",
1003 TaskScope::ApiOverview => "api_overview",
1004 TaskScope::RetrievalQa => "retrieval_qa",
1005 TaskScope::AgentHistory => "agent_history",
1006 }
1007}
1008
1009#[cfg(test)]
1010mod tests {
1011 use super::*;
1012 use crate::budget::CompressionPolicy;
1013
1014 struct MockEstimator(usize);
1015
1016 impl TokenEstimator for MockEstimator {
1017 fn info(&self) -> crate::report::EstimatorInfo {
1018 crate::report::EstimatorInfo {
1019 backend: "mock".to_string(),
1020 model: Some("mock-1".to_string()),
1021 is_exact: true,
1022 }
1023 }
1024
1025 fn count_bytes(&self, _bytes: &[u8]) -> usize {
1026 self.0
1027 }
1028 }
1029
1030 #[test]
1031 fn compress_with_estimator_accepts_a_mock_backend() {
1032 let input = CompressionInput::plain_text(b"hello".to_vec());
1033 let policy = CompressionPolicy::builder().build().unwrap();
1034 let output = compress_with_estimator(input, &policy, &MockEstimator(42)).unwrap();
1035 assert_eq!(output.report.original_tokens, 42);
1036 assert_eq!(output.report.estimator.backend, "mock");
1037 }
1038
1039 #[test]
1040 fn passthrough_when_input_is_already_under_target() {
1041 let input = CompressionInput::plain_text(b"hi".to_vec());
1042 let policy = CompressionPolicy::builder()
1043 .target_tokens(1_000)
1044 .build()
1045 .unwrap();
1046 let output = compress_with_estimator(input.clone(), &policy, &MockEstimator(5)).unwrap();
1047 assert_eq!(output.report.status, Status::Passthrough);
1048 assert_eq!(output.bytes, input.bytes);
1049 }
1050
1051 #[test]
1052 fn unreachable_target_returns_best_effort_bytes_and_never_panics() {
1053 let payload = serde_json::json!({
1054 "messages": [{"role": "system", "content": "a fairly long system prompt here"}]
1055 });
1056 let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
1057 let policy = CompressionPolicy::builder()
1058 .target_tokens(1)
1059 .build()
1060 .unwrap();
1061 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1062
1063 assert_eq!(output.report.status, Status::UnreachableTarget);
1064 assert!(!output.bytes.is_empty());
1065 let budget = output.report.budget.expect("budget report populated");
1066 assert_eq!(budget.target_tokens, Some(1));
1067 assert!(budget.protected_floor > 1);
1068 assert_eq!(budget.achieved_tokens, output.report.compressed_tokens);
1069 }
1070
1071 #[test]
1072 fn no_target_set_runs_pipeline_and_reports_estimator_provenance() {
1073 let input = CompressionInput::plain_text(b"no target here".to_vec());
1074 let policy = CompressionPolicy::builder().build().unwrap();
1075 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1076 assert_eq!(output.report.estimator.backend, "heuristic");
1077 assert_eq!(output.report.status, Status::BestEffort);
1078 }
1079
1080 #[test]
1081 fn public_compress_seam_never_panics_on_empty_input() {
1082 let input = CompressionInput::plain_text(Vec::new());
1083 let policy = CompressionPolicy::builder().build().unwrap();
1084 let output = compress(input, &policy).unwrap();
1085 assert_eq!(output.report.original_tokens, 0);
1086 }
1087
1088 #[test]
1089 fn json_minify_actually_applies_and_reduces_tokens_for_openai_json() {
1090 let payload =
1091 b"{\n \"messages\": [\n {\"role\": \"user\", \"content\": \"hi\"}\n ]\n}".to_vec();
1092 let input = CompressionInput::openai_json(payload);
1093 let policy = CompressionPolicy::builder().build().unwrap();
1094 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1095
1096 let applied = output
1097 .report
1098 .transforms
1099 .iter()
1100 .find(|t| t.id == "json_minify")
1101 .expect("json_minify report present");
1102 assert_eq!(applied.status, TransformStatus::Applied);
1103 assert!(applied.saved_tokens > 0 || applied.tokens_after <= applied.tokens_before);
1104 assert!(serde_json::from_slice::<serde_json::Value>(&output.bytes).is_ok());
1105 }
1106
1107 #[test]
1108 fn json_data_transforms_never_regress_token_count() {
1109 let cases: &[&[u8]] = &[
1113 br#"[{"a":1,"b":2},{"a":3,"b":4},{"a":5,"b":6}]"#, br#"[{"a":1},{"b":2},{"c":3}]"#, br#"{"x":[1,2,3],"y":"already compact scalar"}"#, br#"[1,2,3,4,5,6,7,8,9,10]"#, br#"{}"#, ];
1119 for case in cases {
1120 let input = CompressionInput::json(case.to_vec());
1121 let policy = CompressionPolicy::builder().build().unwrap();
1122 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1123 assert!(
1124 output.report.compressed_tokens <= output.report.original_tokens,
1125 "regressed on {}: {} -> {}",
1126 String::from_utf8_lossy(case),
1127 output.report.original_tokens,
1128 output.report.compressed_tokens,
1129 );
1130 assert!(serde_json::from_slice::<serde_json::Value>(&output.bytes).is_ok());
1133 }
1134 }
1135
1136 #[test]
1137 fn secret_redaction_warning_always_present_when_redaction_runs() {
1138 let input = CompressionInput::plain_text(b"nothing secret here".to_vec());
1139 let policy = CompressionPolicy::builder().build().unwrap();
1140 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1141 assert!(
1142 output
1143 .report
1144 .warnings
1145 .iter()
1146 .any(|w| w.code == WarningCode::UnredactedContentPossible)
1147 );
1148 }
1149
1150 #[test]
1151 fn secret_redaction_removes_a_fake_bearer_token_before_any_other_transform() {
1152 let input = CompressionInput::plain_text(
1153 b"Authorization: Bearer sk-abcdEFGH1234567890123456\nother text".to_vec(),
1154 );
1155 let policy = CompressionPolicy::builder().build().unwrap();
1156 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1157 assert!(!contains(&output.bytes, b"sk-abcdEFGH1234567890123456"));
1158 }
1159
1160 #[test]
1161 fn log_compaction_applies_by_default_after_promotion() {
1162 let mut text = String::from("Starting server on port 8080\n");
1168 for _ in 0..10 {
1169 text.push_str("Connecting to database...\n");
1170 }
1171 text.push_str("Database connection established");
1172 let input = CompressionInput::command_output(text.into_bytes());
1173 let policy = CompressionPolicy::builder()
1174 .task_scope(TaskScope::General)
1175 .build()
1176 .unwrap();
1177 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1178 assert!(
1179 output
1180 .report
1181 .transforms
1182 .iter()
1183 .any(|t| t.id == "log_compaction" && t.status == TransformStatus::Applied)
1184 );
1185 }
1186
1187 #[test]
1188 fn diff_compaction_never_applies_without_experimental_flag() {
1189 let input = CompressionInput::plain_text(
1190 b"diff --git a/f.rs b/f.rs\n@@ -1,2 +1,2 @@\n-old\n+new\n context\n context\n context"
1191 .to_vec(),
1192 );
1193 let policy = CompressionPolicy::builder()
1194 .task_scope(TaskScope::CodeReview)
1195 .build()
1196 .unwrap();
1197 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1198 assert!(
1199 !output
1200 .report
1201 .transforms
1202 .iter()
1203 .any(|t| t.id == "diff_compaction" && t.status == TransformStatus::Applied)
1204 );
1205 }
1206
1207 fn contains(haystack: &[u8], needle: &[u8]) -> bool {
1208 haystack.windows(needle.len()).any(|w| w == needle)
1209 }
1210
1211 static RETRIEVAL_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
1215
1216 fn lock_retrieval_env() -> std::sync::MutexGuard<'static, ()> {
1217 RETRIEVAL_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
1218 }
1219
1220 #[test]
1221 fn store_originals_false_leaves_retrieval_report_absent() {
1222 let input = CompressionInput::plain_text(b"anything".to_vec());
1223 let policy = CompressionPolicy::builder().build().unwrap();
1224 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1225 assert!(output.report.retrieval.is_none());
1226 }
1227
1228 #[test]
1229 fn store_originals_persists_full_payload_and_populates_retrieval_report() {
1230 let _g = lock_retrieval_env();
1231 let dir = std::env::temp_dir().join(format!(
1232 "tokenfold_pipeline_test_store_originals_{}",
1233 std::process::id()
1234 ));
1235 unsafe {
1236 std::env::set_var("XDG_DATA_HOME", &dir);
1237 }
1238
1239 let input = CompressionInput::plain_text(b"nothing secret in here at all".to_vec());
1240 let policy = CompressionPolicy::builder()
1241 .store_originals(true)
1242 .retrieval_namespace("pipeline-test")
1243 .build()
1244 .unwrap();
1245 let output =
1246 compress_with_estimator(input.clone(), &policy, &ByteHeuristicEstimator).unwrap();
1247
1248 let retrieval = output.report.retrieval.expect("retrieval report populated");
1249 assert_eq!(retrieval.marker_count, 1);
1250 assert_eq!(retrieval.persisted_original_bytes, input.bytes.len());
1251 assert_eq!(retrieval.skipped_original_bytes, 0);
1252 assert_eq!(retrieval.store_namespace, "pipeline-test");
1253 assert_eq!(retrieval.hash_algorithm, "sha256");
1254 assert_eq!(
1255 retrieval.ttl_seconds,
1256 Some(crate::retrieval_store::DEFAULT_TTL_SECONDS)
1257 );
1258
1259 let hash = crate::retrieval_store::hex_sha256(&input.bytes);
1260 let store = crate::retrieval_store::RetrievalStore::default_filesystem();
1261 assert_eq!(
1262 store.retrieve(&hash, "pipeline-test"),
1263 crate::retrieval_store::RetrievalOutcome::Found(input.bytes.clone())
1264 );
1265
1266 unsafe {
1267 std::env::remove_var("XDG_DATA_HOME");
1268 }
1269 std::fs::remove_dir_all(&dir).ok();
1270 }
1271
1272 #[test]
1273 fn omitting_lossy_reproduces_todays_output_byte_for_byte() {
1274 let payload = serde_json::json!({"items": (0..12).map(|i| serde_json::json!({"n": i})).collect::<Vec<_>>()});
1278 let bytes = serde_json::to_vec(&payload).unwrap();
1279 let input = CompressionInput::json(bytes);
1280 let policy = CompressionPolicy::builder().build().unwrap();
1281 assert_eq!(policy.lossy, None);
1282 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1283 assert!(
1284 !output
1285 .report
1286 .transforms
1287 .iter()
1288 .any(|t| t.id == "json_prune"),
1289 "json_prune must not even appear in the report when lossy is unset"
1290 );
1291 }
1292
1293 #[test]
1294 fn lossy_prunes_a_large_array_and_reports_json_prune_applied() {
1295 let _g = lock_retrieval_env();
1296 let dir = std::env::temp_dir().join(format!(
1297 "tokenfold_pipeline_test_lossy_{}",
1298 std::process::id()
1299 ));
1300 unsafe {
1301 std::env::set_var("XDG_DATA_HOME", &dir);
1302 }
1303
1304 let padding = "x".repeat(150);
1312 let payload = serde_json::json!({
1313 "items": (0..30).map(|i| serde_json::json!({"id": i, "note": padding})).collect::<Vec<_>>()
1314 });
1315 let bytes = serde_json::to_vec(&payload).unwrap();
1316 let input = CompressionInput::json(bytes);
1317 let policy = CompressionPolicy::builder()
1318 .disable("json_field_fold")
1319 .disable("json_value_dict")
1320 .lossy(crate::budget::LossyPath::Heuristic)
1321 .lossy_ratio(0.2)
1322 .build()
1323 .unwrap();
1324 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1325
1326 let report = output
1327 .report
1328 .transforms
1329 .iter()
1330 .find(|t| t.id == "json_prune")
1331 .expect("json_prune report present");
1332 assert_eq!(report.status, TransformStatus::Applied);
1333 assert!(report.tokens_after < report.tokens_before);
1334 assert!(serde_json::from_slice::<serde_json::Value>(&output.bytes).is_ok());
1335
1336 let out_value: serde_json::Value = serde_json::from_slice(&output.bytes).unwrap();
1338 let arr = out_value["items"].as_array().unwrap();
1339 let store = crate::retrieval_store::RetrievalStore::default_filesystem();
1340 for item in arr {
1341 if let Some(hash) = item
1342 .get("$tf_ref")
1343 .and_then(|r| r.get("hash"))
1344 .and_then(|h| h.as_str())
1345 {
1346 assert!(matches!(
1347 store.retrieve(hash, "default"),
1348 crate::retrieval_store::RetrievalOutcome::Found(_)
1349 ));
1350 }
1351 }
1352
1353 unsafe {
1354 std::env::remove_var("XDG_DATA_HOME");
1355 }
1356 std::fs::remove_dir_all(&dir).ok();
1357 }
1358
1359 #[test]
1360 fn lossy_that_prunes_nothing_is_never_worse_than_a_plain_lossless_run() {
1361 let _g = lock_retrieval_env();
1362 let dir = std::env::temp_dir().join(format!(
1363 "tokenfold_pipeline_test_lossy_noop_{}",
1364 std::process::id()
1365 ));
1366 unsafe {
1367 std::env::set_var("XDG_DATA_HOME", &dir);
1368 }
1369
1370 let payload = serde_json::json!({
1378 "events": (0..12)
1379 .map(|i| serde_json::json!({"seq": i, "retries": 0, "note": "queue drain cycle completed normally"}))
1380 .collect::<Vec<_>>()
1381 });
1382 let bytes = serde_json::to_vec(&payload).unwrap();
1383
1384 let lossless = compress_with_estimator(
1385 CompressionInput::json(bytes.clone()),
1386 &CompressionPolicy::builder().build().unwrap(),
1387 &ByteHeuristicEstimator,
1388 )
1389 .unwrap();
1390 let lossy = compress_with_estimator(
1391 CompressionInput::json(bytes),
1392 &CompressionPolicy::builder()
1393 .lossy(crate::budget::LossyPath::Heuristic)
1394 .lossy_ratio(0.25)
1395 .build()
1396 .unwrap(),
1397 &ByteHeuristicEstimator,
1398 )
1399 .unwrap();
1400
1401 let prune = lossy
1402 .report
1403 .transforms
1404 .iter()
1405 .find(|t| t.id == "json_prune")
1406 .expect("json_prune report present");
1407 assert_ne!(
1408 prune.status,
1409 TransformStatus::Applied,
1410 "fixture must stay in the nothing-was-pruned regime for this test to mean anything"
1411 );
1412 assert_eq!(
1413 lossy.bytes, lossless.bytes,
1414 "a --lossy run that pruned nothing must fall back to the exact lossless output"
1415 );
1416 assert!(
1419 lossy
1420 .report
1421 .transforms
1422 .iter()
1423 .any(|t| t.id == "json_field_fold" && t.status == TransformStatus::Applied),
1424 "the deferred lossless transforms must be replayed once pruning turns out to be a no-op"
1425 );
1426 assert!(lossy.report.quality.is_none());
1428
1429 unsafe {
1430 std::env::remove_var("XDG_DATA_HOME");
1431 }
1432 std::fs::remove_dir_all(&dir).ok();
1433 }
1434
1435 #[test]
1436 fn a_successful_prune_that_folding_would_have_beaten_is_rolled_back() {
1437 let _g = lock_retrieval_env();
1438 let dir = std::env::temp_dir().join(format!(
1439 "tokenfold_pipeline_test_lossy_loses_{}",
1440 std::process::id()
1441 ));
1442 unsafe {
1443 std::env::set_var("XDG_DATA_HOME", &dir);
1444 }
1445
1446 let note = "routine health check completed without incident, all subsystems reported \
1453 green, disk and memory pressure nominal, no retries were required, the \
1454 scheduler handed off cleanly and downstream consumers acknowledged";
1455 let payload = serde_json::json!({
1456 "tasks": (0..30)
1457 .map(|i| serde_json::json!({"id": i, "worker": format!("w-{i}"), "note": note}))
1458 .collect::<Vec<_>>()
1459 });
1460 let bytes = serde_json::to_vec(&payload).unwrap();
1461
1462 let lossless = compress_with_estimator(
1463 CompressionInput::json(bytes.clone()),
1464 &CompressionPolicy::builder().build().unwrap(),
1465 &ByteHeuristicEstimator,
1466 )
1467 .unwrap();
1468 let lossy = compress_with_estimator(
1469 CompressionInput::json(bytes),
1470 &CompressionPolicy::builder()
1471 .lossy(crate::budget::LossyPath::Heuristic)
1472 .lossy_ratio(0.25)
1473 .build()
1474 .unwrap(),
1475 &ByteHeuristicEstimator,
1476 )
1477 .unwrap();
1478
1479 assert_eq!(
1480 lossy.bytes, lossless.bytes,
1481 "a prune that loses to folding must be rolled back in favor of the lossless output"
1482 );
1483 let prune = lossy
1484 .report
1485 .transforms
1486 .iter()
1487 .find(|t| t.id == "json_prune")
1488 .expect("json_prune report present");
1489 assert_ne!(prune.status, TransformStatus::Applied);
1490 assert!(!String::from_utf8_lossy(&lossy.bytes).contains("$tf_ref"));
1496 assert!(lossy.report.quality.is_none());
1497
1498 unsafe {
1499 std::env::remove_var("XDG_DATA_HOME");
1500 }
1501 std::fs::remove_dir_all(&dir).ok();
1502 }
1503
1504 #[test]
1505 fn lossy_is_skipped_when_the_lossless_pipeline_already_meets_the_target() {
1506 let _g = lock_retrieval_env();
1507 let dir = std::env::temp_dir().join(format!(
1508 "tokenfold_pipeline_test_lossy_target_met_{}",
1509 std::process::id()
1510 ));
1511 unsafe {
1512 std::env::set_var("XDG_DATA_HOME", &dir);
1513 }
1514
1515 let padding = "x".repeat(150);
1520 let payload = serde_json::json!({
1521 "items": (0..30).map(|i| serde_json::json!({"id": i, "note": padding})).collect::<Vec<_>>()
1522 });
1523 let bytes = serde_json::to_vec(&payload).unwrap();
1524
1525 let lossless = compress_with_estimator(
1527 CompressionInput::json(bytes.clone()),
1528 &CompressionPolicy::builder().build().unwrap(),
1529 &ByteHeuristicEstimator,
1530 )
1531 .unwrap();
1532 let target = lossless.report.compressed_tokens;
1533
1534 let lossy = compress_with_estimator(
1535 CompressionInput::json(bytes),
1536 &CompressionPolicy::builder()
1537 .target_tokens(target)
1538 .lossy(crate::budget::LossyPath::Heuristic)
1539 .lossy_ratio(0.1)
1540 .build()
1541 .unwrap(),
1542 &ByteHeuristicEstimator,
1543 )
1544 .unwrap();
1545
1546 let prune = lossy
1547 .report
1548 .transforms
1549 .iter()
1550 .find(|t| t.id == "json_prune")
1551 .expect("json_prune report present");
1552 assert_eq!(prune.status, TransformStatus::Skipped);
1553 assert_eq!(prune.skipped_reason, Some(SkippedReason::TargetAlreadyMet));
1554 assert!(!String::from_utf8_lossy(&lossy.bytes).contains("$tf_ref"));
1555 assert_eq!(lossy.report.status, Status::Compressed);
1556
1557 unsafe {
1558 std::env::remove_var("XDG_DATA_HOME");
1559 }
1560 std::fs::remove_dir_all(&dir).ok();
1561 }
1562
1563 #[test]
1564 fn a_rejected_lossy_branch_leaves_no_orphaned_blobs_on_disk() {
1565 let _g = lock_retrieval_env();
1566 let dir = std::env::temp_dir().join(format!(
1567 "tokenfold_pipeline_test_lossy_orphans_{}",
1568 std::process::id()
1569 ));
1570 std::fs::remove_dir_all(&dir).ok();
1571 let store_dir = dir.join("store");
1572
1573 let note = "routine health check completed without incident, all subsystems reported \
1580 green, disk and memory pressure nominal, no retries were required, the \
1581 scheduler handed off cleanly and downstream consumers acknowledged";
1582 let payload = serde_json::json!({
1583 "tasks": (0..30)
1584 .map(|i| serde_json::json!({"id": i, "worker": format!("w-{i}"), "note": note}))
1585 .collect::<Vec<_>>()
1586 });
1587 let policy = CompressionPolicy::builder()
1588 .lossy(crate::budget::LossyPath::Heuristic)
1589 .lossy_ratio(0.25)
1590 .retrieval_store_path(Some(store_dir.clone()))
1591 .build()
1592 .unwrap();
1593 let output = compress_with_estimator(
1594 CompressionInput::json(serde_json::to_vec(&payload).unwrap()),
1595 &policy,
1596 &ByteHeuristicEstimator,
1597 )
1598 .unwrap();
1599
1600 let prune = output
1601 .report
1602 .transforms
1603 .iter()
1604 .find(|t| t.id == "json_prune")
1605 .expect("json_prune report present");
1606 assert_eq!(prune.status, TransformStatus::RolledBack);
1607
1608 let blobs = std::fs::read_dir(store_dir.join("default"))
1611 .map(|rd| {
1612 rd.filter_map(|e| e.ok())
1613 .filter(|e| e.path().extension().is_some_and(|x| x == "bin"))
1614 .count()
1615 })
1616 .unwrap_or(0);
1617 let reported = output
1618 .report
1619 .retrieval
1620 .as_ref()
1621 .map_or(0, |r| r.marker_count);
1622 assert_eq!(
1623 blobs, reported,
1624 "{blobs} blobs on disk vs {reported} reported -- a rejected branch persisted orphans"
1625 );
1626
1627 std::fs::remove_dir_all(&dir).ok();
1628 }
1629
1630 #[test]
1631 fn per_item_stores_report_their_real_ttl_even_when_the_whole_payload_was_refused() {
1632 let _g = lock_retrieval_env();
1633 let dir = std::env::temp_dir().join(format!(
1634 "tokenfold_pipeline_test_lossy_ttl_{}",
1635 std::process::id()
1636 ));
1637 std::fs::remove_dir_all(&dir).ok();
1638 let store_dir = dir.join("store");
1639
1640 let items: Vec<Value> = (0..20)
1650 .map(|i| {
1651 serde_json::json!({
1652 "id": i,
1653 "note": format!("row {i}: {}", (0..40).map(|j| format!("field{}", (i * 7 + j) % 23)).collect::<Vec<_>>().join(" ")),
1654 })
1655 })
1656 .collect();
1657 let payload = serde_json::json!({
1658 "api_key": format!("sk-{}", "A".repeat(40)),
1659 "items": items,
1660 });
1661 let policy = CompressionPolicy::builder()
1662 .lossy(crate::budget::LossyPath::Heuristic)
1663 .lossy_ratio(0.1)
1664 .retrieval_store_path(Some(store_dir))
1665 .build()
1666 .unwrap();
1667 let output = compress_with_estimator(
1668 CompressionInput::json(serde_json::to_vec(&payload).unwrap()),
1669 &policy,
1670 &ByteHeuristicEstimator,
1671 )
1672 .unwrap();
1673
1674 let retrieval = output.report.retrieval.expect("retrieval report present");
1675 assert!(
1676 retrieval.skipped_original_bytes > 0,
1677 "the whole-payload receipt must have been refused for this test to mean anything"
1678 );
1679 assert!(
1680 retrieval.marker_count > 0,
1681 "per-item stores must have succeeded"
1682 );
1683 assert_eq!(
1684 retrieval.ttl_seconds,
1685 Some(retrieval_store::DEFAULT_TTL_SECONDS),
1686 "a positive marker_count must report the TTL those entries were really stored with"
1687 );
1688
1689 std::fs::remove_dir_all(&dir).ok();
1690 }
1691
1692 #[test]
1693 fn a_real_lossy_prune_reports_quality_as_present_but_unvalidated() {
1694 let _g = lock_retrieval_env();
1695 let dir = std::env::temp_dir().join(format!(
1696 "tokenfold_pipeline_test_lossy_quality_{}",
1697 std::process::id()
1698 ));
1699 unsafe {
1700 std::env::set_var("XDG_DATA_HOME", &dir);
1701 }
1702
1703 let padding = "x".repeat(150);
1708 let payload = serde_json::json!({
1709 "items": (0..30).map(|i| serde_json::json!({"id": i, "note": padding})).collect::<Vec<_>>()
1710 });
1711 let policy = CompressionPolicy::builder()
1712 .disable("json_field_fold")
1713 .disable("json_value_dict")
1714 .lossy(crate::budget::LossyPath::Heuristic)
1715 .lossy_ratio(0.2)
1716 .build()
1717 .unwrap();
1718 let output = compress_with_estimator(
1719 CompressionInput::json(serde_json::to_vec(&payload).unwrap()),
1720 &policy,
1721 &ByteHeuristicEstimator,
1722 )
1723 .unwrap();
1724
1725 let quality = output
1726 .report
1727 .quality
1728 .expect("quality present after a lossy run");
1729 assert!(!quality.gate_passed);
1730 assert_eq!(quality.validated_ratio_band, None);
1731 assert_eq!(quality.quality_retention, None);
1732 assert_eq!(quality.contrastive_failure_rate, None);
1733
1734 unsafe {
1735 std::env::remove_var("XDG_DATA_HOME");
1736 }
1737 std::fs::remove_dir_all(&dir).ok();
1738 }
1739
1740 #[test]
1741 fn lossy_on_a_format_it_cannot_run_on_never_persists_the_input() {
1742 let _g = lock_retrieval_env();
1743 let dir = std::env::temp_dir().join(format!(
1744 "tokenfold_pipeline_test_lossy_wrong_format_{}",
1745 std::process::id()
1746 ));
1747 std::fs::remove_dir_all(&dir).ok();
1748 unsafe {
1749 std::env::set_var("XDG_DATA_HOME", &dir);
1750 }
1751
1752 let payload = serde_json::json!({
1757 "model": "gpt-4o",
1758 "messages": (0..8)
1759 .map(|i| serde_json::json!({"role": "user", "content": format!("question {i} {}", "lorem ipsum dolor sit amet ".repeat(20))}))
1760 .collect::<Vec<_>>()
1761 });
1762 let policy = CompressionPolicy::builder()
1763 .lossy(crate::budget::LossyPath::Heuristic)
1764 .lossy_ratio(0.2)
1765 .build()
1766 .unwrap();
1767 let output = compress_with_estimator(
1768 CompressionInput {
1769 format: InputFormat::OpenAiJson,
1770 bytes: serde_json::to_vec(&payload).unwrap(),
1771 },
1772 &policy,
1773 &ByteHeuristicEstimator,
1774 )
1775 .unwrap();
1776
1777 assert_eq!(
1778 output.report.retrieval, None,
1779 "nothing may be persisted when the lossy stage cannot run and --store-originals \
1780 was never asked for"
1781 );
1782 assert!(
1783 !dir.exists(),
1784 "no retrieval directory may be created either"
1785 );
1786
1787 unsafe {
1788 std::env::remove_var("XDG_DATA_HOME");
1789 }
1790 std::fs::remove_dir_all(&dir).ok();
1791 }
1792
1793 #[test]
1794 fn lossy_never_drops_a_secret_shaped_item_fail_closed() {
1795 let _g = lock_retrieval_env();
1796 let dir = std::env::temp_dir().join(format!(
1797 "tokenfold_pipeline_test_lossy_secret_{}",
1798 std::process::id()
1799 ));
1800 unsafe {
1801 std::env::set_var("XDG_DATA_HOME", &dir);
1802 }
1803
1804 let padding = "y".repeat(150);
1808 let payload = serde_json::json!({
1809 "items": (0..10).map(|_| serde_json::json!({
1810 "key": "AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE",
1811 "note": padding,
1812 })).collect::<Vec<_>>()
1813 });
1814 let bytes = serde_json::to_vec(&payload).unwrap();
1815 let input = CompressionInput::json(bytes);
1816 let policy = CompressionPolicy::builder()
1817 .disable("json_field_fold")
1818 .disable("json_value_dict")
1819 .lossy(crate::budget::LossyPath::Heuristic)
1820 .lossy_ratio(0.0)
1821 .unsafe_disable_redaction(true) .build()
1823 .unwrap();
1824 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1825 let out_value: serde_json::Value = serde_json::from_slice(&output.bytes).unwrap();
1826 let arr = out_value["items"].as_array().unwrap();
1827 assert_eq!(arr.len(), 10);
1828 assert!(
1829 arr.iter().all(|item| item.get("$tf_ref").is_none()),
1830 "every item must survive fail-closed since none could be stored"
1831 );
1832
1833 unsafe {
1834 std::env::remove_var("XDG_DATA_HOME");
1835 }
1836 std::fs::remove_dir_all(&dir).ok();
1837 }
1838
1839 #[test]
1840 fn lossy_rollback_never_reports_retrieval_markers_absent_from_the_output() {
1841 struct JointPenaltyEstimator;
1851 impl TokenEstimator for JointPenaltyEstimator {
1852 fn info(&self) -> crate::report::EstimatorInfo {
1853 crate::report::EstimatorInfo {
1854 backend: "test-joint-penalty".to_string(),
1855 model: None,
1856 is_exact: true,
1857 }
1858 }
1859 fn count_bytes(&self, bytes: &[u8]) -> usize {
1860 let has_marker = bytes.windows(7).any(|w| w == b"$tf_ref");
1861 if has_marker && bytes.len() > 300 {
1862 bytes.len() * 3
1863 } else {
1864 bytes.len()
1865 }
1866 }
1867 }
1868
1869 let _g = lock_retrieval_env();
1870 let dir = std::env::temp_dir().join(format!(
1871 "tokenfold_pipeline_test_lossy_rollback_report_{}",
1872 std::process::id()
1873 ));
1874 unsafe {
1875 std::env::set_var("XDG_DATA_HOME", &dir);
1876 }
1877
1878 let padding = "x".repeat(150);
1879 let payload = serde_json::json!({
1880 "items": (0..10).map(|i| serde_json::json!({"n": i, "pad": padding})).collect::<Vec<_>>()
1881 });
1882 let bytes = serde_json::to_vec(&payload).unwrap();
1883 let input = CompressionInput::json(bytes.clone());
1884 let policy = CompressionPolicy::builder()
1885 .disable("json_field_fold")
1886 .disable("json_value_dict")
1887 .disable("json_minify")
1888 .lossy(crate::budget::LossyPath::Heuristic)
1889 .lossy_ratio(0.3)
1890 .build()
1891 .unwrap();
1892 let output = compress_with_estimator(input, &policy, &JointPenaltyEstimator).unwrap();
1893
1894 let jp = output
1895 .report
1896 .transforms
1897 .iter()
1898 .find(|t| t.id == "json_prune")
1899 .expect("json_prune report present");
1900 assert_eq!(
1901 jp.status,
1902 TransformStatus::RolledBack,
1903 "test fixture must actually trigger a rollback to exercise the fix; got {:?}",
1904 jp.status
1905 );
1906 let s = String::from_utf8_lossy(&output.bytes);
1913 assert!(
1914 !s.contains("$tf_ref"),
1915 "rolled-back output must contain no markers"
1916 );
1917 let r = output
1918 .report
1919 .retrieval
1920 .as_ref()
1921 .expect("F-045 whole-payload backup is forced on whenever lossy is set");
1922 assert_eq!(
1923 r.marker_count, 1,
1924 "only the F-045 whole-payload marker should be reported, not any of json_prune's own"
1925 );
1926 assert_eq!(
1927 r.persisted_original_bytes,
1928 bytes.len(),
1929 "persisted bytes must be exactly the whole-payload backup, nothing extra from \
1930 json_prune's own rolled-back per-item drops"
1931 );
1932
1933 unsafe {
1934 std::env::remove_var("XDG_DATA_HOME");
1935 }
1936 std::fs::remove_dir_all(&dir).ok();
1937 }
1938
1939 #[test]
1940 fn lossy_never_drops_the_system_message_or_latest_user_message_for_openai_format() {
1941 let _g = lock_retrieval_env();
1947 let dir = std::env::temp_dir().join(format!(
1948 "tokenfold_pipeline_test_lossy_protected_openai_{}",
1949 std::process::id()
1950 ));
1951 unsafe {
1952 std::env::set_var("XDG_DATA_HOME", &dir);
1953 }
1954
1955 let padding = "x".repeat(300);
1956 let system_content = format!("SYSTEM_PROMPT_MUST_SURVIVE: {padding}");
1957 let payload = serde_json::json!({
1958 "model": "gpt-4",
1959 "messages": [
1960 {"role": "system", "content": system_content},
1961 {"role": "user", "content": format!("turn 1: {padding}")},
1962 {"role": "assistant", "content": format!("turn 2: {padding}")},
1963 {"role": "user", "content": format!("turn 3: {padding}")},
1964 {"role": "assistant", "content": format!("turn 4: {padding}")},
1965 ]
1966 });
1967 let bytes = serde_json::to_vec(&payload).unwrap();
1968 let input = CompressionInput::openai_json(bytes);
1969 let policy = CompressionPolicy::builder()
1970 .disable("json_field_fold")
1971 .disable("json_value_dict")
1972 .lossy(crate::budget::LossyPath::Heuristic)
1973 .lossy_ratio(0.05) .build()
1975 .unwrap();
1976 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
1977
1978 let s = String::from_utf8_lossy(&output.bytes);
1979 assert!(
1980 s.contains("SYSTEM_PROMPT_MUST_SURVIVE"),
1981 "system message must survive lossy pruning byte-for-byte; output was: {s}"
1982 );
1983
1984 let jp = output
1987 .report
1988 .transforms
1989 .iter()
1990 .find(|t| t.id == "json_prune")
1991 .expect("json_prune report present");
1992 assert_eq!(jp.status, TransformStatus::Skipped);
1993 assert_eq!(
1994 jp.skipped_reason,
1995 Some(SkippedReason::NotApplicableToFormat)
1996 );
1997
1998 unsafe {
1999 std::env::remove_var("XDG_DATA_HOME");
2000 }
2001 std::fs::remove_dir_all(&dir).ok();
2002 }
2003
2004 #[test]
2005 fn lossy_only_ever_runs_for_generic_json_never_openai_or_anthropic_format() {
2006 let _g = lock_retrieval_env();
2021 let dir = std::env::temp_dir().join(format!(
2022 "tokenfold_pipeline_test_lossy_duplicate_content_bypass_{}",
2023 std::process::id()
2024 ));
2025 unsafe {
2026 std::env::set_var("XDG_DATA_HOME", &dir);
2027 }
2028
2029 let padding = "x".repeat(150);
2030 let content_x = format!("content X {padding}");
2031 let content_y = format!("content Y {padding}");
2032 let payload = serde_json::json!({
2033 "model": "gpt-4",
2034 "messages": [
2035 {"role": "assistant", "content": content_x, "success": false},
2036 {"role": "system", "content": content_x},
2037 {"role": "assistant", "content": content_y},
2038 {"role": "user", "content": content_x},
2039 ]
2040 });
2041 let bytes = serde_json::to_vec(&payload).unwrap();
2042 let input = CompressionInput::openai_json(bytes);
2043 let policy = CompressionPolicy::builder()
2044 .disable("json_field_fold")
2045 .disable("json_value_dict")
2046 .lossy(crate::budget::LossyPath::Heuristic)
2047 .lossy_ratio(0.35)
2048 .build()
2049 .unwrap();
2050 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
2051
2052 assert!(
2053 !String::from_utf8_lossy(&output.bytes).contains("$tf_ref"),
2054 "no message may be replaced by a marker on a message-format payload"
2055 );
2056 let jp = output
2057 .report
2058 .transforms
2059 .iter()
2060 .find(|t| t.id == "json_prune")
2061 .expect("json_prune report present");
2062 assert_eq!(jp.status, TransformStatus::Skipped);
2063 assert_eq!(
2064 jp.skipped_reason,
2065 Some(SkippedReason::NotApplicableToFormat)
2066 );
2067
2068 unsafe {
2069 std::env::remove_var("XDG_DATA_HOME");
2070 }
2071 std::fs::remove_dir_all(&dir).ok();
2072 }
2073
2074 #[test]
2075 fn lossy_preview_shows_projected_savings_without_any_real_storage_write() {
2076 let _g = lock_retrieval_env();
2081 let dir = std::env::temp_dir().join(format!(
2082 "tokenfold_pipeline_test_lossy_preview_{}",
2083 std::process::id()
2084 ));
2085 unsafe {
2086 std::env::set_var("XDG_DATA_HOME", &dir);
2087 }
2088 let store_root = dir.join("tokenfold").join("retrieve");
2089
2090 let padding = "x".repeat(300);
2091 let payload = serde_json::json!({
2092 "items": (0..20).map(|i| serde_json::json!({"id": i, "note": padding})).collect::<Vec<_>>()
2093 });
2094 let bytes = serde_json::to_vec(&payload).unwrap();
2095 let input = CompressionInput::json(bytes);
2096 let mut policy = CompressionPolicy::builder()
2097 .disable("json_field_fold")
2098 .disable("json_value_dict")
2099 .lossy(crate::budget::LossyPath::Heuristic)
2100 .lossy_ratio(0.2)
2101 .build()
2102 .unwrap();
2103 policy.preview = true;
2104 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
2105
2106 let jp = output
2107 .report
2108 .transforms
2109 .iter()
2110 .find(|t| t.id == "json_prune")
2111 .expect("json_prune report present");
2112 assert_eq!(
2113 jp.status,
2114 TransformStatus::Applied,
2115 "preview must still show projected savings, not silently no-op"
2116 );
2117 assert!(jp.tokens_after < jp.tokens_before);
2118 assert!(
2119 output.report.retrieval.is_none(),
2120 "a preview must not claim anything was persisted"
2121 );
2122 assert!(
2123 !store_root.exists(),
2124 "a preview must never create the retrieval store directory at all"
2125 );
2126
2127 unsafe {
2128 std::env::remove_var("XDG_DATA_HOME");
2129 }
2130 std::fs::remove_dir_all(&dir).ok();
2131 }
2132
2133 #[test]
2134 fn lossy_preview_puts_back_items_a_real_store_would_refuse_over_an_unsafe_namespace() {
2135 let _g = lock_retrieval_env();
2148 let dir = std::env::temp_dir().join(format!(
2149 "tokenfold_pipeline_test_lossy_preview_namespace_refusal_{}",
2150 std::process::id()
2151 ));
2152 unsafe {
2153 std::env::set_var("XDG_DATA_HOME", &dir);
2154 }
2155
2156 let padding = "x".repeat(300);
2157 let items: Vec<Value> = (0..10)
2158 .map(|i| serde_json::json!({"id": i, "note": "boring", "pad": padding}))
2159 .collect();
2160 let payload = serde_json::json!({"items": items});
2161 let bytes = serde_json::to_vec(&payload).unwrap();
2162 let input = CompressionInput::json(bytes);
2163 let mut policy = CompressionPolicy::builder()
2164 .disable("json_field_fold")
2165 .disable("json_value_dict")
2166 .lossy(crate::budget::LossyPath::Heuristic)
2167 .lossy_ratio(0.0) .retrieval_namespace("../escape") .build()
2170 .unwrap();
2171 policy.preview = true;
2172 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
2173
2174 let s = String::from_utf8_lossy(&output.bytes);
2175 assert!(
2176 !s.contains("$tf_ref"),
2177 "every proposed drop must be put back once probed against a namespace a real \
2178 store() call would refuse; output was: {s}"
2179 );
2180 assert!(
2181 output.report.retrieval.is_none(),
2182 "a preview must still never claim anything was really persisted"
2183 );
2184
2185 unsafe {
2186 std::env::remove_var("XDG_DATA_HOME");
2187 }
2188 std::fs::remove_dir_all(&dir).ok();
2189 }
2190
2191 #[test]
2192 fn lossy_preserve_holds_even_when_the_caller_never_disabled_json_field_fold() {
2193 let _g = lock_retrieval_env();
2204 let dir = std::env::temp_dir().join(format!(
2205 "tokenfold_pipeline_test_lossy_preserve_survives_field_fold_{}",
2206 std::process::id()
2207 ));
2208 unsafe {
2209 std::env::set_var("XDG_DATA_HOME", &dir);
2210 }
2211 let items: Vec<Value> = (0..30)
2215 .map(|i| serde_json::json!({"id": i, "note": "x".repeat(300)}))
2216 .collect();
2217 let payload = serde_json::json!({"items": items});
2218 let bytes = serde_json::to_vec(&payload).unwrap();
2219 let input = CompressionInput::json(bytes);
2220 let policy = CompressionPolicy::builder()
2221 .lossy(crate::budget::LossyPath::Heuristic)
2222 .lossy_ratio(0.1)
2223 .lossy_preserve("items")
2224 .build()
2225 .unwrap();
2226 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
2227 let s = String::from_utf8_lossy(&output.bytes);
2228 assert!(
2229 !s.contains("$tf_ref"),
2230 "preserve must hold; output was: {s}"
2231 );
2232
2233 unsafe {
2234 std::env::remove_var("XDG_DATA_HOME");
2235 }
2236 std::fs::remove_dir_all(&dir).ok();
2237 }
2238
2239 #[test]
2240 fn lossy_requires_a_non_memory_backend() {
2241 let err = CompressionPolicy::builder()
2242 .lossy(crate::budget::LossyPath::Heuristic)
2243 .retrieval_backend("memory")
2244 .build();
2245 assert!(
2246 err.is_err(),
2247 "policy construction itself must reject this combination"
2248 );
2249 }
2250
2251 #[test]
2252 fn compress_rejects_a_hand_mutated_policy_that_bypasses_builder_validation() {
2253 let mut policy = CompressionPolicy::builder().build().unwrap();
2257 policy.lossy_ratio = 5.0; let input = CompressionInput::plain_text(b"hello".to_vec());
2259 let err = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap_err();
2260 assert!(matches!(err, TokenFoldError::ConfigError(_)));
2261 }
2262
2263 #[test]
2264 fn store_originals_skips_secret_bearing_payloads_without_erroring_the_compression() {
2265 let _g = lock_retrieval_env();
2266 let dir = std::env::temp_dir().join(format!(
2267 "tokenfold_pipeline_test_store_originals_secret_{}",
2268 std::process::id()
2269 ));
2270 unsafe {
2271 std::env::set_var("XDG_DATA_HOME", &dir);
2272 }
2273
2274 let input =
2275 CompressionInput::plain_text(b"AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE".to_vec());
2276 let policy = CompressionPolicy::builder()
2277 .store_originals(true)
2278 .retrieval_namespace("pipeline-test")
2279 .build()
2280 .unwrap();
2281 let output =
2282 compress_with_estimator(input.clone(), &policy, &ByteHeuristicEstimator).unwrap();
2283
2284 let retrieval = output.report.retrieval.expect("retrieval report populated");
2285 assert_eq!(retrieval.marker_count, 0);
2286 assert_eq!(retrieval.persisted_original_bytes, 0);
2287 assert_eq!(retrieval.skipped_original_bytes, input.bytes.len());
2288
2289 let hash = crate::retrieval_store::hex_sha256(&input.bytes);
2290 let store = crate::retrieval_store::RetrievalStore::default_filesystem();
2291 assert_eq!(
2292 store.retrieve(&hash, "pipeline-test"),
2293 crate::retrieval_store::RetrievalOutcome::Missing
2294 );
2295
2296 unsafe {
2297 std::env::remove_var("XDG_DATA_HOME");
2298 }
2299 std::fs::remove_dir_all(&dir).ok();
2300 }
2301}