1use crate::budget::{CompressionMode, CompressionPolicy, TaskScope, protected_segments};
2use crate::errors::TokenFoldError;
3use crate::input::{CompressionInput, CompressionOutput, InputFormat};
4use crate::modes::{self, ModeEntry, TransformId};
5use crate::report::{
6 BudgetReport, CompressionReport, RetrievalReport, Severity, SkippedReason, TransformReport,
7 TransformStatus, Warning, WarningCode,
8};
9use crate::retrieval_store::{self, RetrievalStore};
10use crate::safety;
11use crate::status::Status;
12use crate::token_estimator::{ByteHeuristicEstimator, TokenEstimator};
13use crate::transforms;
14
15pub fn compress(
18 input: CompressionInput,
19 policy: &CompressionPolicy,
20) -> Result<CompressionOutput, TokenFoldError> {
21 #[cfg(feature = "tiktoken")]
22 {
23 if let Ok(estimator) = crate::token_estimator::TiktokenEstimator::o200k_base() {
24 return compress_with_estimator(input, policy, &estimator);
25 }
26 }
27 compress_with_estimator(input, policy, &ByteHeuristicEstimator)
28}
29
30pub fn compress_with_estimator(
31 input: CompressionInput,
32 policy: &CompressionPolicy,
33 estimator: &dyn TokenEstimator,
34) -> Result<CompressionOutput, TokenFoldError> {
35 let original_tokens = estimator.count_bytes(&input.bytes);
36 let target = policy.target_tokens;
37 let estimator_info = estimator.info();
38
39 let retrieval = maybe_store_originals(&input.bytes, policy);
42
43 if let Some(t) = target
46 && original_tokens <= t
47 {
48 let mut warnings = Vec::new();
49 if !estimator_info.is_exact {
50 warnings.push(heuristic_budget_warning());
51 }
52 let mut report = CompressionReport::new(
53 original_tokens,
54 original_tokens,
55 estimator_info,
56 Status::Passthrough,
57 mode_label(policy.mode).to_string(),
58 format_label(input.format).to_string(),
59 task_scope_label(policy.task_scope).to_string(),
60 Vec::new(),
61 warnings,
62 );
63 report.retrieval = retrieval;
64 return Ok(CompressionOutput {
65 bytes: input.bytes,
66 report,
67 });
68 }
69
70 apply_transforms(input, policy, estimator, original_tokens, target, retrieval)
71}
72
73fn maybe_store_originals(
80 input_bytes: &[u8],
81 policy: &CompressionPolicy,
82) -> Option<RetrievalReport> {
83 if !policy.store_originals {
84 return None;
85 }
86 let ttl_seconds = policy
87 .retrieval_ttl_seconds
88 .unwrap_or(retrieval_store::DEFAULT_TTL_SECONDS);
89 let skipped = || RetrievalReport {
90 store_namespace: policy.retrieval_namespace.clone(),
91 hash_algorithm: "sha256".to_string(),
92 marker_count: 0,
93 ttl_seconds: None,
94 persisted_original_bytes: 0,
95 skipped_original_bytes: input_bytes.len(),
96 };
97 let Ok(store) = RetrievalStore::open(
98 &policy.retrieval_backend,
99 "sha256",
100 policy.retrieval_store_path.clone(),
101 ) else {
102 return Some(skipped());
103 };
104 Some(
105 match store.store(input_bytes, &policy.retrieval_namespace, Some(ttl_seconds)) {
106 Ok(_marker) => RetrievalReport {
107 store_namespace: policy.retrieval_namespace.clone(),
108 hash_algorithm: "sha256".to_string(),
109 marker_count: 1,
110 ttl_seconds: Some(ttl_seconds),
111 persisted_original_bytes: input_bytes.len(),
112 skipped_original_bytes: 0,
113 },
114 Err(_) => skipped(),
115 },
116 )
117}
118
119fn apply_transforms(
120 input: CompressionInput,
121 policy: &CompressionPolicy,
122 estimator: &dyn TokenEstimator,
123 original_tokens: usize,
124 target: Option<usize>,
125 retrieval: Option<RetrievalReport>,
126) -> Result<CompressionOutput, TokenFoldError> {
127 let estimator_info = estimator.info();
128 let mut warnings = Vec::new();
129 let mut transform_reports = Vec::new();
130 if !estimator_info.is_exact {
131 warnings.push(heuristic_budget_warning());
132 }
133
134 let mut bytes;
138 if policy.unsafe_disable_redaction {
139 bytes = input.bytes.clone();
140 warnings.push(Warning {
141 code: WarningCode::UnredactedContentPossible,
142 severity: Severity::Critical,
143 transform: Some("secret_redaction".to_string()),
144 message: "redaction was disabled via unsafe_disable_redaction; output may contain unredacted secrets".to_string(),
145 });
146 transform_reports.push(skipped_at(
147 "secret_redaction",
148 "1.0.0",
149 original_tokens,
150 SkippedReason::DisabledByUser,
151 ));
152 } else {
153 let outcome = transforms::redaction::redact(&input.bytes);
154 let tokens_after = estimator.count_bytes(&outcome.bytes);
155 warnings.push(Warning {
156 code: WarningCode::UnredactedContentPossible,
157 severity: Severity::Info,
158 transform: Some("secret_redaction".to_string()),
159 message: "redaction is best-effort; it is not a guarantee that no secret survives"
160 .to_string(),
161 });
162 transform_reports.push(TransformReport {
163 id: "secret_redaction".to_string(),
164 version: "1.0.0".to_string(),
165 tokens_before: original_tokens,
166 tokens_after,
167 saved_tokens: original_tokens.saturating_sub(tokens_after),
168 savings_ratio: ratio(original_tokens, tokens_after),
169 elapsed_micros: None,
170 status: if outcome.redacted_count > 0 {
171 TransformStatus::Applied
172 } else {
173 TransformStatus::NoOp
174 },
175 skipped_reason: None,
176 warnings: Vec::new(),
177 });
178 bytes = outcome.bytes;
179 }
180
181 let working_input = CompressionInput {
185 format: input.format,
186 bytes: bytes.clone(),
187 };
188 let protected = protected_segments(&working_input, policy);
189 let floor = estimator.count_bytes(&protected.concat());
190 let mut current_tokens = estimator.count_bytes(&bytes);
191
192 if let Some(t) = target
193 && t < floor
194 {
195 warnings.push(Warning {
196 code: WarningCode::UnreachableTarget,
197 severity: Severity::Warn,
198 transform: None,
199 message: format!("target {t} tokens is below the protected floor of {floor} tokens"),
200 });
201 let mut report = CompressionReport::new(
202 original_tokens,
203 current_tokens,
204 estimator_info,
205 Status::UnreachableTarget,
206 mode_label(policy.mode).to_string(),
207 format_label(input.format).to_string(),
208 task_scope_label(policy.task_scope).to_string(),
209 transform_reports,
210 warnings,
211 );
212 report.budget = Some(BudgetReport {
213 target_tokens: target,
214 protected_floor: floor,
215 achieved_tokens: current_tokens,
216 });
217 report.retrieval = retrieval;
218 return Ok(CompressionOutput { bytes, report });
219 }
220
221 let entries = modes::pipeline_for(
224 policy.mode,
225 policy.task_scope,
226 input.format,
227 policy.experimental,
228 &policy.enable,
229 &policy.disabled,
230 );
231 for entry in entries {
232 if let Some(t) = target
233 && current_tokens <= t
234 {
235 transform_reports.push(skipped(
236 entry,
237 current_tokens,
238 SkippedReason::TargetAlreadyMet,
239 ));
240 continue;
241 }
242
243 let tokens_before = current_tokens;
244 let before_bytes = bytes.clone();
245 let max_ratio = entry.max_ratio_for(policy.mode);
246
247 let candidate = match apply_single_transform(entry.transform_id, &bytes, policy) {
248 Ok(candidate) => candidate,
249 Err(_) => {
250 transform_reports.push(skipped(
251 entry,
252 tokens_before,
253 SkippedReason::NotApplicableToFormat,
254 ));
255 continue;
256 }
257 };
258
259 let tokens_after_candidate = estimator.count_bytes(&candidate);
260 if tokens_after_candidate > tokens_before {
261 transform_reports.push(skipped(
263 entry,
264 tokens_before,
265 SkippedReason::WouldIncreaseTokens,
266 ));
267 continue;
268 }
269 if tokens_after_candidate == tokens_before {
270 transform_reports.push(TransformReport {
273 id: entry.transform_id.as_str().to_string(),
274 version: entry.version.to_string(),
275 tokens_before,
276 tokens_after: tokens_before,
277 saved_tokens: 0,
278 savings_ratio: 0.0,
279 elapsed_micros: None,
280 status: TransformStatus::NoOp,
281 skipped_reason: None,
282 warnings: Vec::new(),
283 });
284 continue;
285 }
286 let ratio_used = 1.0 - (tokens_after_candidate as f64 / tokens_before.max(1) as f64);
287 if ratio_used > max_ratio {
288 transform_reports.push(skipped(
289 entry,
290 tokens_before,
291 SkippedReason::NotEnabledInMode,
292 ));
293 continue;
294 }
295
296 if !validate_safety(
297 entry.transform_id,
298 input.format,
299 &before_bytes,
300 &candidate,
301 &protected,
302 ) {
303 transform_reports.push(rolled_back(entry, tokens_before));
304 warnings.push(safety_downgrade_warning(entry.transform_id.as_str()));
305 continue;
306 }
307
308 bytes = candidate;
309 current_tokens = tokens_after_candidate;
310 transform_reports.push(TransformReport {
311 id: entry.transform_id.as_str().to_string(),
312 version: entry.version.to_string(),
313 tokens_before,
314 tokens_after: current_tokens,
315 saved_tokens: tokens_before.saturating_sub(current_tokens),
316 savings_ratio: ratio(tokens_before, current_tokens),
317 elapsed_micros: None,
318 status: TransformStatus::Applied,
319 skipped_reason: None,
320 warnings: Vec::new(),
321 });
322 }
323
324 let status = match target {
325 None => Status::BestEffort,
326 Some(t) if current_tokens <= t => Status::Compressed,
327 Some(_) => Status::BestEffort,
328 };
329
330 let mut report = CompressionReport::new(
331 original_tokens,
332 current_tokens,
333 estimator_info,
334 status,
335 mode_label(policy.mode).to_string(),
336 format_label(input.format).to_string(),
337 task_scope_label(policy.task_scope).to_string(),
338 transform_reports,
339 warnings,
340 );
341 report.budget = Some(BudgetReport {
342 target_tokens: target,
343 protected_floor: floor,
344 achieved_tokens: current_tokens,
345 });
346 report.retrieval = retrieval;
347 Ok(CompressionOutput { bytes, report })
348}
349
350fn apply_single_transform(
351 transform_id: TransformId,
352 bytes: &[u8],
353 policy: &CompressionPolicy,
354) -> Result<Vec<u8>, String> {
355 match transform_id {
356 TransformId::JsonMinify => transforms::json::minify_json(bytes).map_err(|e| e.to_string()),
357 TransformId::JsonFieldFold => {
358 transforms::json_fold::fold_json(bytes).map_err(|e| e.to_string())
359 }
360 TransformId::JsonValueDict => {
361 transforms::json_dict::dict_json(bytes).map_err(|e| e.to_string())
362 }
363 TransformId::SchemaCompaction => {
364 transforms::schema::compact_schema(bytes, 1).map_err(|e| e.to_string())
368 }
369 TransformId::LogCompaction => {
370 let text = std::str::from_utf8(bytes).map_err(|e| e.to_string())?;
371 Ok(transforms::logs::compact(text, false).into_bytes())
372 }
373 TransformId::DiffCompaction => {
374 let text = std::str::from_utf8(bytes).map_err(|e| e.to_string())?;
375 let keep_line_bodies = policy.task_scope != TaskScope::ChangeSummary;
376 Ok(transforms::diff::compact_diff(text, keep_line_bodies).into_bytes())
377 }
378 }
379}
380
381fn validate_safety(
382 transform_id: TransformId,
383 format: InputFormat,
384 before: &[u8],
385 after: &[u8],
386 protected: &[Vec<u8>],
387) -> bool {
388 match transform_id {
389 TransformId::JsonFieldFold => {
393 if !safety::json_still_valid(after) {
394 return false;
395 }
396 if !transforms::json_fold::round_trips(before, after) {
397 return false;
398 }
399 }
400 TransformId::JsonValueDict => {
403 if !safety::json_still_valid(after) {
404 return false;
405 }
406 if !transforms::json_dict::round_trips(before, after) {
407 return false;
408 }
409 }
410 TransformId::JsonMinify | TransformId::SchemaCompaction => {
413 let is_json_format = matches!(
414 format,
415 InputFormat::OpenAiJson | InputFormat::AnthropicJson | InputFormat::Json
416 );
417 if is_json_format {
418 if !safety::json_still_valid(after) {
419 return false;
420 }
421 if !safety::json_key_order_preserved(before, after) {
422 return false;
423 }
424 }
425 }
426 TransformId::LogCompaction | TransformId::DiffCompaction => {}
427 }
428 safety::protected_segments_present(protected, after)
429}
430
431fn skipped(entry: &ModeEntry, tokens: usize, reason: SkippedReason) -> TransformReport {
432 skipped_at(entry.transform_id.as_str(), entry.version, tokens, reason)
433}
434
435fn skipped_at(id: &str, version: &str, tokens: usize, reason: SkippedReason) -> TransformReport {
436 TransformReport {
437 id: id.to_string(),
438 version: version.to_string(),
439 tokens_before: tokens,
440 tokens_after: tokens,
441 saved_tokens: 0,
442 savings_ratio: 0.0,
443 elapsed_micros: None,
444 status: TransformStatus::Skipped,
445 skipped_reason: Some(reason),
446 warnings: Vec::new(),
447 }
448}
449
450fn rolled_back(entry: &ModeEntry, tokens: usize) -> TransformReport {
451 TransformReport {
452 id: entry.transform_id.as_str().to_string(),
453 version: entry.version.to_string(),
454 tokens_before: tokens,
455 tokens_after: tokens,
456 saved_tokens: 0,
457 savings_ratio: 0.0,
458 elapsed_micros: None,
459 status: TransformStatus::RolledBack,
460 skipped_reason: None,
461 warnings: Vec::new(),
462 }
463}
464
465fn safety_downgrade_warning(transform_id: &str) -> Warning {
466 Warning {
467 code: WarningCode::SafetyDowngrade,
468 severity: Severity::Warn,
469 transform: Some(transform_id.to_string()),
470 message: format!(
471 "{transform_id} was rolled back: a safety invariant would have been violated"
472 ),
473 }
474}
475
476fn heuristic_budget_warning() -> Warning {
477 Warning {
478 code: WarningCode::HeuristicBudgetUsed,
479 severity: Severity::Info,
480 transform: None,
481 message: "token counts are heuristic estimates (~bytes/4), not exact".to_string(),
482 }
483}
484
485fn ratio(before: usize, after: usize) -> f64 {
486 if before == 0 {
487 0.0
488 } else {
489 before.saturating_sub(after) as f64 / before as f64
490 }
491}
492
493fn mode_label(mode: CompressionMode) -> &'static str {
494 match mode {
495 CompressionMode::Conservative => "conservative",
496 CompressionMode::Balanced => "balanced",
497 CompressionMode::Aggressive => "aggressive",
498 }
499}
500
501fn format_label(format: InputFormat) -> &'static str {
502 match format {
503 InputFormat::Auto => "auto",
504 InputFormat::OpenAiJson => "openai_json",
505 InputFormat::AnthropicJson => "anthropic_json",
506 InputFormat::Json => "json",
507 InputFormat::PlainText => "plain_text",
508 InputFormat::CommandOutput => "command_output",
509 InputFormat::GitDiff => "git_diff",
510 }
511}
512
513fn task_scope_label(scope: TaskScope) -> &'static str {
514 match scope {
515 TaskScope::All => "all",
516 TaskScope::General => "general",
517 TaskScope::CodeReview => "code_review",
518 TaskScope::ChangeSummary => "change_summary",
519 TaskScope::Debugging => "debugging",
520 TaskScope::Generation => "generation",
521 TaskScope::ApiOverview => "api_overview",
522 TaskScope::RetrievalQa => "retrieval_qa",
523 TaskScope::AgentHistory => "agent_history",
524 }
525}
526
527#[cfg(test)]
528mod tests {
529 use super::*;
530 use crate::budget::CompressionPolicy;
531
532 struct MockEstimator(usize);
533
534 impl TokenEstimator for MockEstimator {
535 fn info(&self) -> crate::report::EstimatorInfo {
536 crate::report::EstimatorInfo {
537 backend: "mock".to_string(),
538 model: Some("mock-1".to_string()),
539 is_exact: true,
540 }
541 }
542
543 fn count_bytes(&self, _bytes: &[u8]) -> usize {
544 self.0
545 }
546 }
547
548 #[test]
549 fn compress_with_estimator_accepts_a_mock_backend() {
550 let input = CompressionInput::plain_text(b"hello".to_vec());
551 let policy = CompressionPolicy::builder().build().unwrap();
552 let output = compress_with_estimator(input, &policy, &MockEstimator(42)).unwrap();
553 assert_eq!(output.report.original_tokens, 42);
554 assert_eq!(output.report.estimator.backend, "mock");
555 }
556
557 #[test]
558 fn passthrough_when_input_is_already_under_target() {
559 let input = CompressionInput::plain_text(b"hi".to_vec());
560 let policy = CompressionPolicy::builder()
561 .target_tokens(1_000)
562 .build()
563 .unwrap();
564 let output = compress_with_estimator(input.clone(), &policy, &MockEstimator(5)).unwrap();
565 assert_eq!(output.report.status, Status::Passthrough);
566 assert_eq!(output.bytes, input.bytes);
567 }
568
569 #[test]
570 fn unreachable_target_returns_best_effort_bytes_and_never_panics() {
571 let payload = serde_json::json!({
572 "messages": [{"role": "system", "content": "a fairly long system prompt here"}]
573 });
574 let input = CompressionInput::openai_json(serde_json::to_vec(&payload).unwrap());
575 let policy = CompressionPolicy::builder()
576 .target_tokens(1)
577 .build()
578 .unwrap();
579 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
580
581 assert_eq!(output.report.status, Status::UnreachableTarget);
582 assert!(!output.bytes.is_empty());
583 let budget = output.report.budget.expect("budget report populated");
584 assert_eq!(budget.target_tokens, Some(1));
585 assert!(budget.protected_floor > 1);
586 assert_eq!(budget.achieved_tokens, output.report.compressed_tokens);
587 }
588
589 #[test]
590 fn no_target_set_runs_pipeline_and_reports_estimator_provenance() {
591 let input = CompressionInput::plain_text(b"no target here".to_vec());
592 let policy = CompressionPolicy::builder().build().unwrap();
593 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
594 assert_eq!(output.report.estimator.backend, "heuristic");
595 assert_eq!(output.report.status, Status::BestEffort);
596 }
597
598 #[test]
599 fn public_compress_seam_never_panics_on_empty_input() {
600 let input = CompressionInput::plain_text(Vec::new());
601 let policy = CompressionPolicy::builder().build().unwrap();
602 let output = compress(input, &policy).unwrap();
603 assert_eq!(output.report.original_tokens, 0);
604 }
605
606 #[test]
607 fn json_minify_actually_applies_and_reduces_tokens_for_openai_json() {
608 let payload =
609 b"{\n \"messages\": [\n {\"role\": \"user\", \"content\": \"hi\"}\n ]\n}".to_vec();
610 let input = CompressionInput::openai_json(payload);
611 let policy = CompressionPolicy::builder().build().unwrap();
612 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
613
614 let applied = output
615 .report
616 .transforms
617 .iter()
618 .find(|t| t.id == "json_minify")
619 .expect("json_minify report present");
620 assert_eq!(applied.status, TransformStatus::Applied);
621 assert!(applied.saved_tokens > 0 || applied.tokens_after <= applied.tokens_before);
622 assert!(serde_json::from_slice::<serde_json::Value>(&output.bytes).is_ok());
623 }
624
625 #[test]
626 fn json_data_transforms_never_regress_token_count() {
627 let cases: &[&[u8]] = &[
631 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#"{}"#, ];
637 for case in cases {
638 let input = CompressionInput::json(case.to_vec());
639 let policy = CompressionPolicy::builder().build().unwrap();
640 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
641 assert!(
642 output.report.compressed_tokens <= output.report.original_tokens,
643 "regressed on {}: {} -> {}",
644 String::from_utf8_lossy(case),
645 output.report.original_tokens,
646 output.report.compressed_tokens,
647 );
648 assert!(serde_json::from_slice::<serde_json::Value>(&output.bytes).is_ok());
651 }
652 }
653
654 #[test]
655 fn secret_redaction_warning_always_present_when_redaction_runs() {
656 let input = CompressionInput::plain_text(b"nothing secret here".to_vec());
657 let policy = CompressionPolicy::builder().build().unwrap();
658 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
659 assert!(
660 output
661 .report
662 .warnings
663 .iter()
664 .any(|w| w.code == WarningCode::UnredactedContentPossible)
665 );
666 }
667
668 #[test]
669 fn secret_redaction_removes_a_fake_bearer_token_before_any_other_transform() {
670 let input = CompressionInput::plain_text(
671 b"Authorization: Bearer sk-abcdEFGH1234567890123456\nother text".to_vec(),
672 );
673 let policy = CompressionPolicy::builder().build().unwrap();
674 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
675 assert!(!contains(&output.bytes, b"sk-abcdEFGH1234567890123456"));
676 }
677
678 #[test]
679 fn log_compaction_applies_by_default_after_promotion() {
680 let mut text = String::from("Starting server on port 8080\n");
686 for _ in 0..10 {
687 text.push_str("Connecting to database...\n");
688 }
689 text.push_str("Database connection established");
690 let input = CompressionInput::command_output(text.into_bytes());
691 let policy = CompressionPolicy::builder()
692 .task_scope(TaskScope::General)
693 .build()
694 .unwrap();
695 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
696 assert!(
697 output
698 .report
699 .transforms
700 .iter()
701 .any(|t| t.id == "log_compaction" && t.status == TransformStatus::Applied)
702 );
703 }
704
705 #[test]
706 fn diff_compaction_never_applies_without_experimental_flag() {
707 let input = CompressionInput::plain_text(
708 b"diff --git a/f.rs b/f.rs\n@@ -1,2 +1,2 @@\n-old\n+new\n context\n context\n context"
709 .to_vec(),
710 );
711 let policy = CompressionPolicy::builder()
712 .task_scope(TaskScope::CodeReview)
713 .build()
714 .unwrap();
715 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
716 assert!(
717 !output
718 .report
719 .transforms
720 .iter()
721 .any(|t| t.id == "diff_compaction" && t.status == TransformStatus::Applied)
722 );
723 }
724
725 fn contains(haystack: &[u8], needle: &[u8]) -> bool {
726 haystack.windows(needle.len()).any(|w| w == needle)
727 }
728
729 static RETRIEVAL_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
733
734 fn lock_retrieval_env() -> std::sync::MutexGuard<'static, ()> {
735 RETRIEVAL_ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner())
736 }
737
738 #[test]
739 fn store_originals_false_leaves_retrieval_report_absent() {
740 let input = CompressionInput::plain_text(b"anything".to_vec());
741 let policy = CompressionPolicy::builder().build().unwrap();
742 let output = compress_with_estimator(input, &policy, &ByteHeuristicEstimator).unwrap();
743 assert!(output.report.retrieval.is_none());
744 }
745
746 #[test]
747 fn store_originals_persists_full_payload_and_populates_retrieval_report() {
748 let _g = lock_retrieval_env();
749 let dir = std::env::temp_dir().join(format!(
750 "tokenfold_pipeline_test_store_originals_{}",
751 std::process::id()
752 ));
753 unsafe {
754 std::env::set_var("XDG_DATA_HOME", &dir);
755 }
756
757 let input = CompressionInput::plain_text(b"nothing secret in here at all".to_vec());
758 let policy = CompressionPolicy::builder()
759 .store_originals(true)
760 .retrieval_namespace("pipeline-test")
761 .build()
762 .unwrap();
763 let output =
764 compress_with_estimator(input.clone(), &policy, &ByteHeuristicEstimator).unwrap();
765
766 let retrieval = output.report.retrieval.expect("retrieval report populated");
767 assert_eq!(retrieval.marker_count, 1);
768 assert_eq!(retrieval.persisted_original_bytes, input.bytes.len());
769 assert_eq!(retrieval.skipped_original_bytes, 0);
770 assert_eq!(retrieval.store_namespace, "pipeline-test");
771 assert_eq!(retrieval.hash_algorithm, "sha256");
772 assert_eq!(
773 retrieval.ttl_seconds,
774 Some(crate::retrieval_store::DEFAULT_TTL_SECONDS)
775 );
776
777 let hash = crate::retrieval_store::hex_sha256(&input.bytes);
778 let store = crate::retrieval_store::RetrievalStore::default_filesystem();
779 assert_eq!(
780 store.retrieve(&hash, "pipeline-test"),
781 crate::retrieval_store::RetrievalOutcome::Found(input.bytes.clone())
782 );
783
784 unsafe {
785 std::env::remove_var("XDG_DATA_HOME");
786 }
787 std::fs::remove_dir_all(&dir).ok();
788 }
789
790 #[test]
791 fn store_originals_skips_secret_bearing_payloads_without_erroring_the_compression() {
792 let _g = lock_retrieval_env();
793 let dir = std::env::temp_dir().join(format!(
794 "tokenfold_pipeline_test_store_originals_secret_{}",
795 std::process::id()
796 ));
797 unsafe {
798 std::env::set_var("XDG_DATA_HOME", &dir);
799 }
800
801 let input =
802 CompressionInput::plain_text(b"AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE".to_vec());
803 let policy = CompressionPolicy::builder()
804 .store_originals(true)
805 .retrieval_namespace("pipeline-test")
806 .build()
807 .unwrap();
808 let output =
809 compress_with_estimator(input.clone(), &policy, &ByteHeuristicEstimator).unwrap();
810
811 let retrieval = output.report.retrieval.expect("retrieval report populated");
812 assert_eq!(retrieval.marker_count, 0);
813 assert_eq!(retrieval.persisted_original_bytes, 0);
814 assert_eq!(retrieval.skipped_original_bytes, input.bytes.len());
815
816 let hash = crate::retrieval_store::hex_sha256(&input.bytes);
817 let store = crate::retrieval_store::RetrievalStore::default_filesystem();
818 assert_eq!(
819 store.retrieve(&hash, "pipeline-test"),
820 crate::retrieval_store::RetrievalOutcome::Missing
821 );
822
823 unsafe {
824 std::env::remove_var("XDG_DATA_HOME");
825 }
826 std::fs::remove_dir_all(&dir).ok();
827 }
828}