1use std::path::Path;
2use std::sync::{Arc, Mutex};
3
4use crate::ast_parser::AstParser;
5use crate::budget_tracker::{BudgetTracker, UsageReport};
6use crate::cache_manager::CacheManager;
7use crate::confidence_router::ConfidenceRouter;
8use crate::cost_calculator::{CostCalculator, SessionCostSummary};
9use crate::ctx_format::CtxFormat;
10use crate::error::{Result, SqzError};
11use crate::model_router::ModelRouter;
12use crate::pin_manager::PinManager;
13use crate::pipeline::CompressionPipeline;
14use crate::plugin_api::PluginLoader;
15use crate::preset::{Preset, PresetParser};
16use crate::session_store::{SessionStore, SessionSummary};
17use crate::terse_mode::TerseMode;
18use crate::types::{CompressedContent, PinEntry, Provenance, SessionId};
19use crate::verifier::Verifier;
20
21pub struct SqzEngine {
32 preset: Arc<Mutex<Preset>>,
34 pipeline: Arc<Mutex<CompressionPipeline>>,
35 model_router: Arc<Mutex<ModelRouter>>,
36
37 session_store: SessionStore,
39 cache_manager: CacheManager,
40 budget_tracker: BudgetTracker,
41 cost_calculator: CostCalculator,
42 ast_parser: AstParser,
43 terse_mode: TerseMode,
44 pin_manager: PinManager,
45 confidence_router: ConfidenceRouter,
46 _plugin_loader: PluginLoader,
47}
48
49impl SqzEngine {
50 pub fn new() -> Result<Self> {
55 let preset = Preset::default();
56 let store_path = Self::default_store_path();
57 Self::with_preset_and_store(preset, &store_path)
58 }
59
60 fn default_store_path() -> std::path::PathBuf {
69 if let Some(custom) = Self::store_path_from_env(std::env::var("SQZ_DB_PATH").ok()) {
70 return custom;
71 }
72 if let Some(home) = dirs_next::home_dir() {
73 let sqz_dir = home.join(".sqz");
74 if std::fs::create_dir_all(&sqz_dir).is_ok() {
75 #[cfg(unix)]
78 {
79 use std::os::unix::fs::PermissionsExt;
80 let _ = std::fs::set_permissions(
81 &sqz_dir,
82 std::fs::Permissions::from_mode(0o700),
83 );
84 }
85 return sqz_dir.join("sessions.db");
86 }
87 }
88 let dir = std::env::temp_dir();
90 dir.join(format!(
91 "sqz_session_{}_{}.db",
92 std::process::id(),
93 std::time::SystemTime::now()
94 .duration_since(std::time::UNIX_EPOCH)
95 .map(|d| d.as_nanos())
96 .unwrap_or(0)
97 ))
98 }
99
100 fn store_path_from_env(value: Option<String>) -> Option<std::path::PathBuf> {
104 let raw = value?;
105 let trimmed = raw.trim();
106 if trimmed.is_empty() {
107 return None;
108 }
109 let path = std::path::PathBuf::from(trimmed);
110 if let Some(parent) = path.parent() {
111 if !parent.as_os_str().is_empty() && !parent.exists() {
112 if std::fs::create_dir_all(parent).is_ok() {
115 #[cfg(unix)]
116 {
117 use std::os::unix::fs::PermissionsExt;
118 let _ = std::fs::set_permissions(
119 parent,
120 std::fs::Permissions::from_mode(0o700),
121 );
122 }
123 }
124 }
125 }
126 Some(path)
127 }
128
129 pub fn with_preset_and_store(preset: Preset, store_path: &Path) -> Result<Self> {
135 let pipeline = CompressionPipeline::new(&preset);
136 let window_size = preset.budget.default_window_size;
137
138 let session_store = SessionStore::open_or_create(store_path)?;
140 let cache_store = SessionStore::open_or_create(store_path)?;
141 let pin_store = SessionStore::open_or_create(store_path)?;
142
143 Ok(SqzEngine {
144 preset: Arc::new(Mutex::new(preset.clone())),
145 pipeline: Arc::new(Mutex::new(pipeline)),
146 model_router: Arc::new(Mutex::new(ModelRouter::new(&preset))),
147 session_store,
148 cache_manager: CacheManager::new(cache_store, 512 * 1024 * 1024),
149 budget_tracker: BudgetTracker::new(window_size, &preset),
150 cost_calculator: CostCalculator::with_defaults(),
151 ast_parser: AstParser::new(),
152 terse_mode: TerseMode,
153 pin_manager: PinManager::new(pin_store),
154 confidence_router: ConfidenceRouter::new(),
155 _plugin_loader: PluginLoader::new(Path::new("plugins")),
156 })
157 }
158
159 pub fn compress(&self, input: &str) -> Result<CompressedContent> {
167 let preset = self.preset.lock()
168 .map_err(|_| SqzError::Other("preset lock poisoned".into()))?;
169 let pipeline = self.pipeline.lock()
170 .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
171 let ctx = crate::pipeline::SessionContext {
172 session_id: "engine".to_string(),
173 };
174
175 let mode = self.confidence_router.route(input);
177
178 if mode == crate::confidence_router::CompressionMode::Safe {
180 if std::env::var_os("SQZ_VERBOSE").is_some() {
183 eprintln!("[sqz] safe mode: high-risk content passed through unchanged");
184 }
185 return self.compress_safe(input, &pipeline, &ctx);
186 }
187
188 let mut result = pipeline.compress(input, &ctx, &preset)?;
190
191 let verify = Verifier::verify(input, &result.data);
193 let fallback = verify.fallback_triggered;
194 result.verify = Some(verify);
195
196 if fallback && result.data != input {
198 eprintln!("[sqz] fallback: verifier confidence {:.2} below threshold — re-compressing in safe mode",
199 result.verify.as_ref().map(|v| v.confidence).unwrap_or(0.0));
200 let safe_result = self.compress_safe(input, &pipeline, &ctx)?;
201 return Ok(safe_result);
202 }
203
204 Ok(result)
205 }
206
207 pub fn compress_with_cache(&self, input: &str) -> Result<crate::cache_manager::CacheResult> {
223 let pipeline = self.pipeline.lock()
224 .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
225 self.cache_manager.get_or_compress(
231 std::path::Path::new(""),
232 input.as_bytes(),
233 &pipeline,
234 )
235 }
236
237 pub fn compress_with_cache_lossless(&self, input: &str) -> Result<crate::cache_manager::CacheResult> {
243 let pipeline = self.pipeline.lock()
244 .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
245 self.cache_manager.get_or_compress_lossless(
246 std::path::Path::new(""),
247 input.as_bytes(),
248 &pipeline,
249 )
250 }
251
252 pub fn compress_or_passthrough(&self, input: &str) -> CompressedContent {
260 match self.compress(input) {
261 Ok(result) => result,
262 Err(_) => {
263 let tokens = (input.len() as u32 + 3) / 4;
264 CompressedContent {
265 data: input.to_string(),
266 tokens_compressed: tokens,
267 tokens_original: tokens,
268 stages_applied: vec![],
269 compression_ratio: 1.0,
270 provenance: crate::types::Provenance::default(),
271 verify: None,
272 }
273 }
274 }
275 }
276
277 pub fn compress_with_mode(&self, input: &str, mode: crate::confidence_router::CompressionMode) -> Result<CompressedContent> {
283 let pipeline = self.pipeline.lock()
284 .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
285 let ctx = crate::pipeline::SessionContext {
286 session_id: "engine".to_string(),
287 };
288
289 match mode {
290 crate::confidence_router::CompressionMode::Safe => {
291 self.compress_safe(input, &pipeline, &ctx)
292 }
293 _ => {
294 drop(pipeline); self.compress(input)
297 }
298 }
299 }
300
301 fn compress_safe(
303 &self,
304 input: &str,
305 pipeline: &crate::pipeline::CompressionPipeline,
306 ctx: &crate::pipeline::SessionContext,
307 ) -> Result<CompressedContent> {
308 use crate::preset::{
309 CompressionConfig, CondenseConfig, CustomTransformsConfig, BudgetConfig,
310 ModelConfig, PresetMeta, TerseModeConfig, TerseLevel, ToolSelectionConfig,
311 };
312
313 let safe_preset = Preset {
314 preset: PresetMeta {
315 name: "safe".to_string(),
316 version: "1.0".to_string(),
317 description: "Safe fallback — minimal compression".to_string(),
318 },
319 compression: CompressionConfig {
320 stages: vec!["condense".to_string()],
321 keep_fields: None,
322 strip_fields: None,
323 condense: Some(CondenseConfig { enabled: true, max_repeated_lines: 3 }),
324 git_diff_fold: None,
325 strip_nulls: None,
326 flatten: None,
327 truncate_strings: None,
328 collapse_arrays: None,
329 custom_transforms: Some(CustomTransformsConfig { enabled: false }),
330 },
331 tool_selection: ToolSelectionConfig {
332 max_tools: 5,
333 similarity_threshold: 0.7,
334 default_tools: vec![],
335 },
336 budget: BudgetConfig {
337 warning_threshold: 0.70,
338 ceiling_threshold: 0.85,
339 default_window_size: 200_000,
340 agents: Default::default(),
341 },
342 terse_mode: TerseModeConfig { enabled: false, level: TerseLevel::Moderate },
343 model: ModelConfig {
344 family: "anthropic".to_string(),
345 primary: String::new(),
346 local: String::new(),
347 complexity_threshold: 0.4,
348 pricing: None,
349 },
350 };
351
352 let mut result = pipeline.compress(input, ctx, &safe_preset)?;
353 let verify = Verifier::verify(input, &result.data);
354 result.verify = Some(verify);
355 result.provenance = Provenance {
356 label: Some("safe-fallback".to_string()),
357 ..Default::default()
358 };
359 Ok(result)
360 }
361
362 pub fn compress_with_provenance(
364 &self,
365 input: &str,
366 provenance: Provenance,
367 ) -> Result<CompressedContent> {
368 let mut result = self.compress(input)?;
369 result.provenance = provenance;
370 Ok(result)
371 }
372
373 pub fn export_ctx(&self, session_id: &str) -> Result<String> {
375 let session = self.session_store.load_session(session_id.to_string())?;
376 CtxFormat::serialize(&session)
377 }
378
379 pub fn import_ctx(&self, ctx: &str) -> Result<SessionId> {
381 let session = CtxFormat::deserialize(ctx)?;
382 self.session_store.save_session(&session)
383 }
384
385 pub fn pin(&self, session_id: &str, turn_index: usize, reason: &str, tokens: u32) -> Result<PinEntry> {
387 self.pin_manager.pin(session_id, turn_index, reason, tokens)
388 }
389
390 pub fn unpin(&self, session_id: &str, turn_index: usize) -> Result<()> {
392 self.pin_manager.unpin(session_id, turn_index)
393 }
394
395 pub fn search_sessions(&self, query: &str) -> Result<Vec<SessionSummary>> {
397 self.session_store.search(query)
398 }
399
400 pub fn usage_report(&self, agent_id: &str) -> UsageReport {
402 self.budget_tracker.usage_report(agent_id.to_string())
403 }
404
405 pub fn cost_summary(&self, session_id: &str) -> Result<SessionCostSummary> {
407 let session = self.session_store.load_session(session_id.to_string())?;
408 Ok(self.cost_calculator.session_summary(&session))
409 }
410
411 pub fn reload_preset(&mut self, toml: &str) -> Result<()> {
413 let new_preset = PresetParser::parse(toml)?;
414 if let Ok(mut pipeline) = self.pipeline.lock() {
415 pipeline.reload_preset(&new_preset)?;
416 }
417 if let Ok(mut router) = self.model_router.lock() {
418 *router = ModelRouter::new(&new_preset);
419 }
420 if let Ok(mut preset) = self.preset.lock() {
421 *preset = new_preset;
422 }
423 Ok(())
424 }
425
426 pub fn watch_preset_file(&self, path: &Path) -> Result<notify::RecommendedWatcher> {
431 use notify::{Event, EventKind, RecursiveMode, Watcher};
432
433 let preset_arc = Arc::clone(&self.preset);
434 let pipeline_arc = Arc::clone(&self.pipeline);
435 let router_arc = Arc::clone(&self.model_router);
436 let watched_path = path.to_owned();
437
438 let mut watcher = notify::recommended_watcher(move |res: notify::Result<Event>| {
439 if let Ok(event) = res {
440 if matches!(event.kind, EventKind::Modify(_) | EventKind::Create(_)) {
441 match std::fs::read_to_string(&watched_path) {
442 Ok(toml_str) => match PresetParser::parse(&toml_str) {
443 Ok(new_preset) => {
444 if let Ok(mut p) = pipeline_arc.lock() {
445 let _ = p.reload_preset(&new_preset);
446 }
447 if let Ok(mut r) = router_arc.lock() {
448 *r = ModelRouter::new(&new_preset);
449 }
450 if let Ok(mut pr) = preset_arc.lock() {
451 *pr = new_preset;
452 }
453 }
454 Err(e) => eprintln!("[sqz] invalid preset: {e}"),
455 },
456 Err(e) => eprintln!("[sqz] preset read error: {e}"),
457 }
458 }
459 }
460 })
461 .map_err(|e| SqzError::Other(format!("watcher error: {e}")))?;
462
463 watcher
464 .watch(path, RecursiveMode::NonRecursive)
465 .map_err(|e| SqzError::Other(format!("watch error: {e}")))?;
466
467 Ok(watcher)
468 }
469
470 pub fn session_store(&self) -> &SessionStore {
472 &self.session_store
473 }
474
475 pub fn cache_manager(&self) -> &CacheManager {
477 &self.cache_manager
478 }
479
480 pub fn ast_parser(&self) -> &AstParser {
482 &self.ast_parser
483 }
484
485 pub fn terse_mode(&self) -> &TerseMode {
487 &self.terse_mode
488 }
489
490 pub fn reorder_context(
496 &self,
497 sections: &mut Vec<crate::litm_positioner::ContextSection>,
498 strategy: crate::litm_positioner::LitmStrategy,
499 ) {
500 let positioner = crate::litm_positioner::LitmPositioner::new(strategy);
501 positioner.reorder(sections);
502 }
503
504 pub fn route_compression_mode(&self, content: &str) -> crate::confidence_router::CompressionMode {
507 self.confidence_router.route(content)
508 }
509
510 pub fn count_tokens(&self, text: &str) -> u32 {
512 let preset = self.preset.lock().unwrap();
513 self.pipeline.lock().unwrap().count_tokens(text, &preset)
514 }
515}
516
517
518#[cfg(test)]
519mod tests {
520 use super::*;
521 use crate::types::{BudgetState, CorrectionLog, ModelFamily, SessionState};
522 use chrono::Utc;
523 use std::path::PathBuf;
524
525 fn make_session(id: &str) -> SessionState {
526 let now = Utc::now();
527 SessionState {
528 id: id.to_string(),
529 project_dir: PathBuf::from("/tmp/test"),
530 conversation: vec![],
531 corrections: CorrectionLog::default(),
532 pins: vec![],
533 learnings: vec![],
534 compressed_summary: "test session".to_string(),
535 budget: BudgetState {
536 window_size: 200_000,
537 consumed: 0,
538 pinned: 0,
539 model_family: ModelFamily::AnthropicClaude,
540 },
541 tool_usage: vec![],
542 created_at: now,
543 updated_at: now,
544 }
545 }
546
547 #[test]
548 fn test_engine_new() {
549 let engine = SqzEngine::new();
550 assert!(engine.is_ok(), "SqzEngine::new() should succeed");
551 }
552
553 #[test]
554 fn test_compress_or_passthrough_returns_result_on_valid_input() {
555 let engine = SqzEngine::new().unwrap();
556 let result = engine.compress_or_passthrough("hello world");
557 assert_eq!(result.data, "hello world");
558 assert!(result.tokens_original > 0);
559 }
560
561 #[test]
562 fn test_compress_or_passthrough_never_panics_on_empty() {
563 let engine = SqzEngine::new().unwrap();
564 let result = engine.compress_or_passthrough("");
565 assert_eq!(result.data, "");
566 assert_eq!(result.compression_ratio, 1.0);
567 }
568
569 #[test]
570 fn test_compress_or_passthrough_handles_json() {
571 let engine = SqzEngine::new().unwrap();
572 let result = engine.compress_or_passthrough(r#"{"key":"value"}"#);
573 assert!(!result.data.is_empty());
575 }
576
577 #[test]
578 fn test_compress_or_passthrough_handles_binary_garbage() {
579 let engine = SqzEngine::new().unwrap();
580 let garbage = "\x00\x01\x02\x7f invalid control chars \t\n\r";
582 let result = engine.compress_or_passthrough(garbage);
583 assert!(!result.data.is_empty());
584 }
585
586 #[test]
587 fn test_compress_plain_text() {
588 let engine = SqzEngine::new().unwrap();
589 let result = engine.compress("hello world");
590 assert!(result.is_ok());
591 assert_eq!(result.unwrap().data, "hello world");
592 }
593
594 #[test]
595 fn test_compress_json_applies_toon() {
596 let engine = SqzEngine::new().unwrap();
597 let result = engine.compress(r#"{"name":"Alice","age":30}"#).unwrap();
598 assert!(result.data.starts_with("TOON:"), "JSON should be TOON-encoded");
599 }
600
601 #[test]
602 fn test_export_import_ctx_round_trip() {
603 let dir = tempfile::tempdir().unwrap();
604 let store_path = dir.path().join("store.db");
605 let engine = SqzEngine::with_preset_and_store(Preset::default(), &store_path).unwrap();
606
607 let session = make_session("sess-rt");
608 engine.session_store().save_session(&session).unwrap();
609
610 let ctx = engine.export_ctx("sess-rt").unwrap();
611 let imported_id = engine.import_ctx(&ctx).unwrap();
612 assert_eq!(imported_id, "sess-rt");
613 }
614
615 #[test]
616 fn test_search_sessions() {
617 let dir = tempfile::tempdir().unwrap();
618 let store_path = dir.path().join("store.db");
619 let engine = SqzEngine::with_preset_and_store(Preset::default(), &store_path).unwrap();
620
621 let mut session = make_session("sess-search");
622 session.compressed_summary = "authentication refactor".to_string();
623 engine.session_store().save_session(&session).unwrap();
624
625 let results = engine.search_sessions("authentication").unwrap();
626 assert_eq!(results.len(), 1);
627 assert_eq!(results[0].id, "sess-search");
628 }
629
630 #[test]
631 fn test_usage_report_starts_at_zero() {
632 let engine = SqzEngine::new().unwrap();
633 let report = engine.usage_report("default");
634 assert_eq!(report.consumed, 0);
635 assert_eq!(report.available, report.allocated);
636 }
637
638 #[test]
639 fn test_cost_summary() {
640 let dir = tempfile::tempdir().unwrap();
641 let store_path = dir.path().join("store.db");
642 let engine = SqzEngine::with_preset_and_store(Preset::default(), &store_path).unwrap();
643
644 let session = make_session("sess-cost");
645 engine.session_store().save_session(&session).unwrap();
646
647 let summary = engine.cost_summary("sess-cost").unwrap();
648 assert_eq!(summary.total_tokens, 0);
649 assert!((summary.total_usd - 0.0).abs() < f64::EPSILON);
650 }
651
652 #[test]
653 fn test_reload_preset_updates_state() {
654 let mut engine = SqzEngine::new().unwrap();
655 let toml = r#"
656[preset]
657name = "reloaded"
658version = "2.0"
659
660[compression]
661stages = []
662
663[tool_selection]
664max_tools = 5
665similarity_threshold = 0.7
666
667[budget]
668warning_threshold = 0.70
669ceiling_threshold = 0.85
670default_window_size = 200000
671
672[terse_mode]
673enabled = false
674level = "moderate"
675
676[model]
677family = "anthropic"
678primary = "claude-sonnet-4-20250514"
679complexity_threshold = 0.4
680"#;
681 assert!(engine.reload_preset(toml).is_ok());
682 let preset = engine.preset.lock().unwrap();
684 assert_eq!(preset.preset.name, "reloaded");
685 }
686
687 #[test]
688 fn test_reload_invalid_preset_returns_error() {
689 let mut engine = SqzEngine::new().unwrap();
690 let result = engine.reload_preset("not valid toml [[[");
691 assert!(result.is_err(), "invalid TOML should return error");
692 }
693
694 #[test]
695 fn test_export_nonexistent_session_returns_error() {
696 let engine = SqzEngine::new().unwrap();
697 let result = engine.export_ctx("does-not-exist");
698 assert!(result.is_err());
699 }
700
701 #[test]
702 fn test_import_invalid_ctx_returns_error() {
703 let engine = SqzEngine::new().unwrap();
704 let result = engine.import_ctx("not valid json {{{");
705 assert!(result.is_err());
706 }
707
708 #[test]
709 fn store_path_env_override_used_when_set() {
710 let dir = tempfile::tempdir().unwrap();
711 let target = dir.path().join("nested").join("project.db");
712 let resolved =
713 SqzEngine::store_path_from_env(Some(target.to_string_lossy().to_string()));
714 assert_eq!(resolved, Some(target.clone()));
715 assert!(target.parent().unwrap().exists());
717 }
718
719 #[test]
720 fn store_path_env_ignores_empty_values() {
721 assert_eq!(SqzEngine::store_path_from_env(None), None);
722 assert_eq!(SqzEngine::store_path_from_env(Some(String::new())), None);
723 assert_eq!(SqzEngine::store_path_from_env(Some(" ".to_string())), None);
724 }
725
726 #[test]
727 fn store_path_env_leaves_existing_parent_permissions_alone() {
728 let dir = tempfile::tempdir().unwrap();
729 #[cfg(unix)]
730 {
731 use std::os::unix::fs::PermissionsExt;
732 std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755))
733 .unwrap();
734 }
735 let target = dir.path().join("sessions.db");
736 let resolved =
737 SqzEngine::store_path_from_env(Some(target.to_string_lossy().to_string()));
738 assert_eq!(resolved, Some(target));
739 #[cfg(unix)]
740 {
741 use std::os::unix::fs::PermissionsExt;
742 let mode = std::fs::metadata(dir.path()).unwrap().permissions().mode() & 0o777;
743 assert_eq!(mode, 0o755, "pre-existing parent must not be chmodded");
744 }
745 }
746
747 #[test]
748 fn engine_with_custom_store_path_round_trips() {
749 let dir = tempfile::tempdir().unwrap();
750 let db = dir.path().join("proj").join("sessions.db");
751 let resolved = SqzEngine::store_path_from_env(Some(db.to_string_lossy().to_string()))
752 .unwrap();
753 let engine = SqzEngine::with_preset_and_store(Preset::default(), &resolved).unwrap();
754 engine
755 .session_store()
756 .log_compression(100, 40, &[], "test")
757 .unwrap();
758 assert!(db.exists(), "database file must land at the custom path");
759 let stats = engine.session_store().compression_stats().unwrap();
760 assert_eq!(stats.total_compressions, 1);
761 }
762}