1use std::ffi::OsStr;
4use std::io::Read as _;
5use std::path::{Path, PathBuf};
6use std::time::{Duration, Instant};
7
8use agentic_memory::{
9 AmemReader, AmemWriter, CognitiveEventBuilder, Edge, EdgeType, EventType, MemoryGraph,
10 QueryEngine, WriteEngine,
11};
12use serde_json::Value;
13
14use crate::types::{McpError, McpResult};
15
16const DEFAULT_AUTO_SAVE_SECS: u64 = 30;
18const DEFAULT_BACKUP_INTERVAL_SECS: u64 = 900;
20const DEFAULT_BACKUP_RETENTION: usize = 24;
22const DEFAULT_SLEEP_CYCLE_SECS: u64 = 1800;
24const DEFAULT_ARCHIVE_MIN_SESSION_NODES: usize = 25;
26const DEFAULT_HOT_MIN_DECAY: f32 = 0.7;
28const DEFAULT_WARM_MIN_DECAY: f32 = 0.3;
30const DEFAULT_SLA_MAX_MUTATIONS_PER_MIN: u32 = 240;
32const DEFAULT_HEALTH_LEDGER_EMIT_SECS: u64 = 30;
34const DEFAULT_STORAGE_BUDGET_BYTES: u64 = 2 * 1024 * 1024 * 1024;
36const DEFAULT_STORAGE_BUDGET_HORIZON_YEARS: u32 = 20;
38const DEFAULT_AUTO_CAPTURE_MAX_CHARS: usize = 2048;
40const CURRENT_AMEM_VERSION: u32 = 1;
42
43#[derive(Debug, Clone, Copy)]
44enum AutonomicProfile {
45 Desktop,
46 Cloud,
47 Aggressive,
48}
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51enum StorageMigrationPolicy {
52 AutoSafe,
53 Strict,
54 Off,
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58enum StorageBudgetMode {
59 AutoRollup,
60 Warn,
61 Off,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65enum AutoCaptureMode {
66 Safe,
68 Full,
70 Off,
72}
73
74#[derive(Debug, Clone, Copy)]
75struct ProfileDefaults {
76 auto_save_secs: u64,
77 backup_secs: u64,
78 backup_retention: usize,
79 sleep_cycle_secs: u64,
80 sleep_idle_secs: u64,
81 archive_min_session_nodes: usize,
82 hot_min_decay: f32,
83 warm_min_decay: f32,
84 sla_max_mutations_per_min: u32,
85}
86
87impl AutonomicProfile {
88 fn from_env(name: &str) -> Self {
89 let raw = read_env_string(name).unwrap_or_else(|| "desktop".to_string());
90 match raw.trim().to_ascii_lowercase().as_str() {
91 "cloud" => Self::Cloud,
92 "aggressive" => Self::Aggressive,
93 _ => Self::Desktop,
94 }
95 }
96
97 fn defaults(self) -> ProfileDefaults {
98 match self {
99 Self::Desktop => ProfileDefaults {
100 auto_save_secs: DEFAULT_AUTO_SAVE_SECS,
101 backup_secs: DEFAULT_BACKUP_INTERVAL_SECS,
102 backup_retention: DEFAULT_BACKUP_RETENTION,
103 sleep_cycle_secs: DEFAULT_SLEEP_CYCLE_SECS,
104 sleep_idle_secs: 180,
105 archive_min_session_nodes: DEFAULT_ARCHIVE_MIN_SESSION_NODES,
106 hot_min_decay: DEFAULT_HOT_MIN_DECAY,
107 warm_min_decay: DEFAULT_WARM_MIN_DECAY,
108 sla_max_mutations_per_min: DEFAULT_SLA_MAX_MUTATIONS_PER_MIN,
109 },
110 Self::Cloud => ProfileDefaults {
111 auto_save_secs: 15,
112 backup_secs: 600,
113 backup_retention: 48,
114 sleep_cycle_secs: 900,
115 sleep_idle_secs: 90,
116 archive_min_session_nodes: 50,
117 hot_min_decay: 0.75,
118 warm_min_decay: 0.4,
119 sla_max_mutations_per_min: 600,
120 },
121 Self::Aggressive => ProfileDefaults {
122 auto_save_secs: 10,
123 backup_secs: 300,
124 backup_retention: 16,
125 sleep_cycle_secs: 300,
126 sleep_idle_secs: 45,
127 archive_min_session_nodes: 15,
128 hot_min_decay: 0.8,
129 warm_min_decay: 0.5,
130 sla_max_mutations_per_min: 900,
131 },
132 }
133 }
134
135 fn as_str(self) -> &'static str {
136 match self {
137 Self::Desktop => "desktop",
138 Self::Cloud => "cloud",
139 Self::Aggressive => "aggressive",
140 }
141 }
142}
143
144impl StorageMigrationPolicy {
145 fn from_env(name: &str) -> Self {
146 let raw = read_env_string(name).unwrap_or_else(|| "auto-safe".to_string());
147 match raw.trim().to_ascii_lowercase().as_str() {
148 "strict" => Self::Strict,
149 "off" | "disabled" | "none" => Self::Off,
150 _ => Self::AutoSafe,
151 }
152 }
153
154 fn as_str(self) -> &'static str {
155 match self {
156 Self::AutoSafe => "auto-safe",
157 Self::Strict => "strict",
158 Self::Off => "off",
159 }
160 }
161}
162
163impl StorageBudgetMode {
164 fn from_env(name: &str) -> Self {
165 let raw = read_env_string(name).unwrap_or_else(|| "auto-rollup".to_string());
166 match raw.trim().to_ascii_lowercase().as_str() {
167 "warn" => Self::Warn,
168 "off" | "disabled" | "none" => Self::Off,
169 _ => Self::AutoRollup,
170 }
171 }
172
173 fn as_str(self) -> &'static str {
174 match self {
175 Self::AutoRollup => "auto-rollup",
176 Self::Warn => "warn",
177 Self::Off => "off",
178 }
179 }
180}
181
182impl AutoCaptureMode {
183 fn from_env(name: &str) -> Self {
184 let raw = read_env_string(name).unwrap_or_else(|| "safe".to_string());
185 match raw.trim().to_ascii_lowercase().as_str() {
186 "full" => Self::Full,
187 "off" | "disabled" | "none" => Self::Off,
188 _ => Self::Safe,
189 }
190 }
191
192 fn as_str(self) -> &'static str {
193 match self {
194 Self::Safe => "safe",
195 Self::Full => "full",
196 Self::Off => "off",
197 }
198 }
199}
200
201pub struct SessionManager {
203 graph: MemoryGraph,
204 query_engine: QueryEngine,
205 write_engine: WriteEngine,
206 file_path: PathBuf,
207 current_session: u32,
208 profile: AutonomicProfile,
209 migration_policy: StorageMigrationPolicy,
210 dirty: bool,
211 last_save: Instant,
212 auto_save_interval: Duration,
213 backup_interval: Duration,
214 backup_retention: usize,
215 backups_dir: PathBuf,
216 save_generation: u64,
217 last_backup_generation: u64,
218 last_backup: Instant,
219 sleep_cycle_interval: Duration,
220 archive_min_session_nodes: usize,
221 hot_min_decay: f32,
222 warm_min_decay: f32,
223 sla_max_mutations_per_min: u32,
224 last_sleep_cycle: Instant,
225 sleep_idle_min: Duration,
226 last_activity: Instant,
227 mutation_window_started: Instant,
228 mutation_window_count: u32,
229 maintenance_throttle_count: u64,
230 last_health_ledger_emit: Instant,
231 health_ledger_emit_interval: Duration,
232 storage_budget_mode: StorageBudgetMode,
233 storage_budget_max_bytes: u64,
234 storage_budget_horizon_years: u32,
235 storage_budget_target_fraction: f32,
236 storage_budget_rollup_count: u64,
237 auto_capture_mode: AutoCaptureMode,
238 auto_capture_redact: bool,
239 auto_capture_max_chars: usize,
240 auto_capture_count: u64,
241}
242
243impl SessionManager {
244 pub fn open(path: &str) -> McpResult<Self> {
246 let file_path = PathBuf::from(path);
247 let dimension = agentic_memory::DEFAULT_DIMENSION;
248 let file_existed = file_path.exists();
249 let profile = AutonomicProfile::from_env("AMEM_AUTONOMIC_PROFILE");
250 let defaults = profile.defaults();
251 let migration_policy = StorageMigrationPolicy::from_env("AMEM_STORAGE_MIGRATION_POLICY");
252 let detected_version = if file_existed {
253 read_storage_version(&file_path)
254 } else {
255 None
256 };
257 let legacy_version = detected_version.filter(|v| *v < CURRENT_AMEM_VERSION);
258
259 let graph = if file_existed {
260 tracing::info!("Opening existing memory file: {}", file_path.display());
261 AmemReader::read_from_file(&file_path)
262 .map_err(|e| McpError::AgenticMemory(format!("Failed to read memory file: {e}")))?
263 } else {
264 tracing::info!("Creating new memory file: {}", file_path.display());
265 if let Some(parent) = file_path.parent() {
267 std::fs::create_dir_all(parent).map_err(|e| {
268 McpError::Io(std::io::Error::other(format!(
269 "Failed to create directory {}: {e}",
270 parent.display()
271 )))
272 })?;
273 }
274 MemoryGraph::new(dimension)
275 };
276
277 let session_ids = graph.session_index().session_ids();
279 let current_session = session_ids.iter().copied().max().unwrap_or(0) + 1;
280
281 tracing::info!(
282 "Session {} started. Graph has {} nodes, {} edges.",
283 current_session,
284 graph.node_count(),
285 graph.edge_count()
286 );
287 tracing::info!(
288 "Autonomic profile={} migration_policy={}",
289 profile.as_str(),
290 migration_policy.as_str()
291 );
292
293 let auto_save_secs = read_env_u64("AMEM_AUTOSAVE_SECS", defaults.auto_save_secs);
294 let backup_secs = read_env_u64("AMEM_AUTO_BACKUP_SECS", defaults.backup_secs).max(30);
295 let backup_retention =
296 read_env_usize("AMEM_AUTO_BACKUP_RETENTION", defaults.backup_retention).max(1);
297 let backups_dir = resolve_backups_dir(&file_path);
298 let sleep_cycle_secs =
299 read_env_u64("AMEM_SLEEP_CYCLE_SECS", defaults.sleep_cycle_secs).max(60);
300 let sleep_idle_secs =
301 read_env_u64("AMEM_SLEEP_IDLE_SECS", defaults.sleep_idle_secs).max(30);
302 let archive_min_session_nodes = read_env_usize(
303 "AMEM_ARCHIVE_MIN_SESSION_NODES",
304 defaults.archive_min_session_nodes,
305 )
306 .max(1);
307 let hot_min_decay =
308 read_env_f32("AMEM_TIER_HOT_MIN_DECAY", defaults.hot_min_decay).clamp(0.0, 1.0);
309 let warm_min_decay = read_env_f32("AMEM_TIER_WARM_MIN_DECAY", defaults.warm_min_decay)
310 .clamp(0.0, 1.0)
311 .min(hot_min_decay);
312 let sla_max_mutations_per_min = read_env_u32(
313 "AMEM_SLA_MAX_MUTATIONS_PER_MIN",
314 defaults.sla_max_mutations_per_min,
315 )
316 .max(1);
317 let health_ledger_emit_interval = Duration::from_secs(
318 read_env_u64(
319 "AMEM_HEALTH_LEDGER_EMIT_SECS",
320 DEFAULT_HEALTH_LEDGER_EMIT_SECS,
321 )
322 .max(5),
323 );
324 let storage_budget_mode = StorageBudgetMode::from_env("AMEM_STORAGE_BUDGET_MODE");
325 let storage_budget_max_bytes =
326 read_env_u64("AMEM_STORAGE_BUDGET_BYTES", DEFAULT_STORAGE_BUDGET_BYTES).max(1);
327 let storage_budget_horizon_years = read_env_u32(
328 "AMEM_STORAGE_BUDGET_HORIZON_YEARS",
329 DEFAULT_STORAGE_BUDGET_HORIZON_YEARS,
330 )
331 .max(1);
332 let storage_budget_target_fraction =
333 read_env_f32("AMEM_STORAGE_BUDGET_TARGET_FRACTION", 0.85).clamp(0.50, 0.99);
334 let auto_capture_mode = AutoCaptureMode::from_env("AMEM_AUTO_CAPTURE_MODE");
335 let auto_capture_redact = read_env_bool("AMEM_AUTO_CAPTURE_REDACT", true);
336 let auto_capture_max_chars = read_env_usize(
337 "AMEM_AUTO_CAPTURE_MAX_CHARS",
338 DEFAULT_AUTO_CAPTURE_MAX_CHARS,
339 )
340 .clamp(256, 16384);
341
342 let mut manager = Self {
343 graph,
344 query_engine: QueryEngine::new(),
345 write_engine: WriteEngine::new(dimension),
346 file_path,
347 current_session,
348 profile,
349 migration_policy,
350 dirty: false,
351 last_save: Instant::now(),
352 auto_save_interval: Duration::from_secs(auto_save_secs),
353 backup_interval: Duration::from_secs(backup_secs),
354 backup_retention,
355 backups_dir,
356 save_generation: if file_existed { 1 } else { 0 },
357 last_backup_generation: 0,
358 last_backup: Instant::now(),
359 sleep_cycle_interval: Duration::from_secs(sleep_cycle_secs),
360 archive_min_session_nodes,
361 hot_min_decay,
362 warm_min_decay,
363 sla_max_mutations_per_min,
364 last_sleep_cycle: Instant::now(),
365 sleep_idle_min: Duration::from_secs(sleep_idle_secs),
366 last_activity: Instant::now(),
367 mutation_window_started: Instant::now(),
368 mutation_window_count: 0,
369 maintenance_throttle_count: 0,
370 last_health_ledger_emit: Instant::now()
371 .checked_sub(health_ledger_emit_interval)
372 .unwrap_or_else(Instant::now),
373 health_ledger_emit_interval,
374 storage_budget_mode,
375 storage_budget_max_bytes,
376 storage_budget_horizon_years,
377 storage_budget_target_fraction,
378 storage_budget_rollup_count: 0,
379 auto_capture_mode,
380 auto_capture_redact,
381 auto_capture_max_chars,
382 auto_capture_count: 0,
383 };
384
385 if let Some(version) = legacy_version {
386 match migration_policy {
387 StorageMigrationPolicy::Strict => {
388 return Err(McpError::AgenticMemory(format!(
389 "Legacy .amem version {} blocked by strict migration policy",
390 version
391 )));
392 }
393 StorageMigrationPolicy::Off => {
394 tracing::warn!(
395 "Legacy storage version detected (v{}), auto-migration disabled by policy",
396 version
397 );
398 }
399 StorageMigrationPolicy::AutoSafe => {
400 if let Some(checkpoint) = manager.create_migration_checkpoint(version)? {
401 tracing::info!(
402 "Legacy storage version detected (v{}), checkpoint created at {}",
403 version,
404 checkpoint.display()
405 );
406 }
407 manager.dirty = true;
408 manager.save()?;
409 tracing::info!(
410 "Auto-migrated memory storage from v{} to v{} at {}",
411 version,
412 CURRENT_AMEM_VERSION,
413 manager.file_path.display()
414 );
415 }
416 }
417 }
418
419 Ok(manager)
420 }
421
422 pub fn graph(&self) -> &MemoryGraph {
424 &self.graph
425 }
426
427 pub fn graph_mut(&mut self) -> &mut MemoryGraph {
429 self.dirty = true;
430 self.last_activity = Instant::now();
431 self.record_mutation();
432 &mut self.graph
433 }
434
435 pub fn query_engine(&self) -> &QueryEngine {
437 &self.query_engine
438 }
439
440 pub fn write_engine(&self) -> &WriteEngine {
442 &self.write_engine
443 }
444
445 pub fn current_session_id(&self) -> u32 {
447 self.current_session
448 }
449
450 pub fn start_session(&mut self, explicit_id: Option<u32>) -> McpResult<u32> {
452 let session_id = explicit_id.unwrap_or_else(|| {
453 let ids = self.graph.session_index().session_ids();
454 ids.iter().copied().max().unwrap_or(0) + 1
455 });
456
457 self.current_session = session_id;
458 self.last_activity = Instant::now();
459 tracing::info!("Started session {session_id}");
460 Ok(session_id)
461 }
462
463 pub fn end_session_with_episode(&mut self, session_id: u32, summary: &str) -> McpResult<u64> {
465 let episode_id = self
466 .write_engine
467 .compress_session(&mut self.graph, session_id, summary)
468 .map_err(|e| McpError::AgenticMemory(format!("Failed to compress session: {e}")))?;
469
470 self.dirty = true;
471 self.last_activity = Instant::now();
472 self.record_mutation();
473 self.save()?;
474
475 tracing::info!("Ended session {session_id}, created episode node {episode_id}");
476
477 Ok(episode_id)
478 }
479
480 pub fn save(&mut self) -> McpResult<()> {
482 if !self.dirty {
483 return Ok(());
484 }
485
486 let writer = AmemWriter::new(self.graph.dimension());
487 writer
488 .write_to_file(&self.graph, &self.file_path)
489 .map_err(|e| McpError::AgenticMemory(format!("Failed to write memory file: {e}")))?;
490
491 self.dirty = false;
492 self.last_save = Instant::now();
493 self.save_generation = self.save_generation.saturating_add(1);
494 tracing::debug!("Saved memory file: {}", self.file_path.display());
495 Ok(())
496 }
497
498 pub fn maybe_auto_save(&mut self) -> McpResult<()> {
500 if self.dirty && self.last_save.elapsed() >= self.auto_save_interval {
501 self.save()?;
502 }
503 Ok(())
504 }
505
506 pub fn run_maintenance_tick(&mut self) -> McpResult<()> {
508 if self.should_throttle_maintenance() {
509 self.maintenance_throttle_count = self.maintenance_throttle_count.saturating_add(1);
510 self.maybe_auto_save()?;
511 self.emit_health_ledger("throttled")?;
512 tracing::debug!(
513 "Maintenance throttled by SLA guard: mutation_rate={} threshold={}",
514 self.mutation_rate_per_min(),
515 self.sla_max_mutations_per_min
516 );
517 return Ok(());
518 }
519
520 self.maybe_run_sleep_cycle()?;
521 self.maybe_auto_save()?;
522 self.maybe_enforce_storage_budget()?;
523 self.maybe_auto_backup()?;
524 self.emit_health_ledger("normal")?;
525 Ok(())
526 }
527
528 pub fn maybe_run_sleep_cycle(&mut self) -> McpResult<()> {
530 if self.last_sleep_cycle.elapsed() < self.sleep_cycle_interval {
531 return Ok(());
532 }
533 if self.last_activity.elapsed() < self.sleep_idle_min {
534 return Ok(());
535 }
536
537 let now = agentic_memory::now_micros();
538 let decay_report = self
539 .write_engine
540 .run_decay(&mut self.graph, now)
541 .map_err(|e| McpError::AgenticMemory(format!("Sleep-cycle decay failed: {e}")))?;
542 let archived_sessions = self.auto_archive_completed_sessions()?;
543
544 if decay_report.nodes_decayed > 0 || archived_sessions > 0 {
545 self.dirty = true;
546 self.save()?;
547 }
548
549 let (hot, warm, cold) = self.tier_counts();
550 self.last_sleep_cycle = Instant::now();
551 tracing::info!(
552 "Sleep-cycle complete: decayed={} archived_sessions={} tiers(h/w/c)={}/{}/{}",
553 decay_report.nodes_decayed,
554 archived_sessions,
555 hot,
556 warm,
557 cold
558 );
559 Ok(())
560 }
561
562 pub fn maybe_auto_backup(&mut self) -> McpResult<()> {
564 if self.last_backup.elapsed() < self.backup_interval {
565 return Ok(());
566 }
567 if self.save_generation <= self.last_backup_generation {
568 return Ok(());
569 }
570 if !self.file_path.exists() {
571 return Ok(());
572 }
573
574 std::fs::create_dir_all(&self.backups_dir).map_err(McpError::Io)?;
575 let backup_path = self.next_backup_path();
576 std::fs::copy(&self.file_path, &backup_path).map_err(McpError::Io)?;
577 self.last_backup_generation = self.save_generation;
578 self.last_backup = Instant::now();
579 self.prune_old_backups()?;
580 tracing::info!("Auto-backup written: {}", backup_path.display());
581 Ok(())
582 }
583
584 pub fn mark_dirty(&mut self) {
586 self.dirty = true;
587 self.last_activity = Instant::now();
588 self.record_mutation();
589 }
590
591 pub fn file_path(&self) -> &PathBuf {
593 &self.file_path
594 }
595
596 pub fn maintenance_interval(&self) -> Duration {
598 self.auto_save_interval
599 .min(self.backup_interval)
600 .min(self.sleep_cycle_interval)
601 }
602
603 pub fn capture_prompt_request(
605 &mut self,
606 prompt_name: &str,
607 arguments: Option<&Value>,
608 ) -> McpResult<Option<u64>> {
609 if self.auto_capture_mode == AutoCaptureMode::Off {
610 return Ok(None);
611 }
612 match extract_prompt_capture_text(prompt_name, arguments)? {
613 Some(text) => self.persist_auto_capture(EventType::Fact, &text, 0.90),
614 None => Ok(None),
615 }
616 }
617
618 pub fn capture_tool_call(
620 &mut self,
621 tool_name: &str,
622 arguments: Option<&Value>,
623 ) -> McpResult<Option<u64>> {
624 if self.auto_capture_mode == AutoCaptureMode::Off {
625 return Ok(None);
626 }
627
628 let text = match self.auto_capture_mode {
629 AutoCaptureMode::Safe => extract_safe_tool_capture_text(tool_name, arguments)?,
630 AutoCaptureMode::Full => extract_full_tool_capture_text(tool_name, arguments)?,
631 AutoCaptureMode::Off => None,
632 };
633 match text {
634 Some(v) => self.persist_auto_capture(EventType::Inference, &v, 0.82),
635 None => Ok(None),
636 }
637 }
638
639 pub fn add_event(
641 &mut self,
642 event_type: EventType,
643 content: &str,
644 confidence: f32,
645 edges: Vec<(u64, EdgeType, f32)>,
646 ) -> McpResult<(u64, usize)> {
647 let event = CognitiveEventBuilder::new(event_type, content.to_string())
648 .session_id(self.current_session)
649 .confidence(confidence)
650 .build();
651
652 let result = self
654 .write_engine
655 .ingest(&mut self.graph, vec![event], vec![])
656 .map_err(|e| McpError::AgenticMemory(format!("Failed to add event: {e}")))?;
657
658 let node_id = result.new_node_ids.first().copied().ok_or_else(|| {
659 McpError::InternalError("No node ID returned from ingest".to_string())
660 })?;
661
662 let mut edge_count = 0;
664 for (target_id, edge_type, weight) in &edges {
665 let edge = Edge::new(node_id, *target_id, *edge_type, *weight);
666 self.graph
667 .add_edge(edge)
668 .map_err(|e| McpError::AgenticMemory(format!("Failed to add edge: {e}")))?;
669 edge_count += 1;
670 }
671
672 self.dirty = true;
673 self.last_activity = Instant::now();
674 self.record_mutation();
675 self.maybe_auto_save()?;
676
677 Ok((node_id, edge_count))
678 }
679
680 pub fn correct_node(&mut self, old_node_id: u64, new_content: &str) -> McpResult<u64> {
682 let new_id = self
683 .write_engine
684 .correct(
685 &mut self.graph,
686 old_node_id,
687 new_content,
688 self.current_session,
689 )
690 .map_err(|e| McpError::AgenticMemory(format!("Failed to correct node: {e}")))?;
691
692 self.dirty = true;
693 self.last_activity = Instant::now();
694 self.record_mutation();
695 self.maybe_auto_save()?;
696
697 Ok(new_id)
698 }
699
700 fn record_mutation(&mut self) {
701 if self.mutation_window_started.elapsed() >= Duration::from_secs(60) {
702 self.mutation_window_started = Instant::now();
703 self.mutation_window_count = 0;
704 }
705 self.mutation_window_count = self.mutation_window_count.saturating_add(1);
706 }
707
708 fn mutation_rate_per_min(&self) -> u32 {
709 let elapsed = self.mutation_window_started.elapsed().as_secs().max(1);
710 let scaled = (self.mutation_window_count as u64)
711 .saturating_mul(60)
712 .saturating_div(elapsed);
713 scaled.min(u32::MAX as u64) as u32
714 }
715
716 fn should_throttle_maintenance(&self) -> bool {
717 self.mutation_rate_per_min() > self.sla_max_mutations_per_min
718 }
719
720 fn emit_health_ledger(&mut self, maintenance_mode: &str) -> McpResult<()> {
721 if self.last_health_ledger_emit.elapsed() < self.health_ledger_emit_interval {
722 return Ok(());
723 }
724
725 let dir = resolve_health_ledger_dir();
726 std::fs::create_dir_all(&dir).map_err(McpError::Io)?;
727 let path = dir.join("agentic-memory.json");
728 let tmp = dir.join("agentic-memory.json.tmp");
729 let (hot, warm, cold) = self.tier_counts();
730 let current_size_bytes = self.current_file_size_bytes();
731 let projected_size_bytes = self.projected_file_size_bytes(current_size_bytes);
732 let over_budget = current_size_bytes > self.storage_budget_max_bytes
733 || projected_size_bytes
734 .map(|v| v > self.storage_budget_max_bytes)
735 .unwrap_or(false);
736 let payload = serde_json::json!({
737 "project": "AgenticMemory",
738 "timestamp": chrono::Utc::now().to_rfc3339(),
739 "status": "ok",
740 "autonomic": {
741 "profile": self.profile.as_str(),
742 "migration_policy": self.migration_policy.as_str(),
743 "maintenance_mode": maintenance_mode,
744 "throttle_count": self.maintenance_throttle_count,
745 },
746 "sla": {
747 "mutation_rate_per_min": self.mutation_rate_per_min(),
748 "max_mutations_per_min": self.sla_max_mutations_per_min
749 },
750 "storage": {
751 "file": self.file_path.display().to_string(),
752 "dirty": self.dirty,
753 "save_generation": self.save_generation,
754 "backup_retention": self.backup_retention,
755 },
756 "storage_budget": {
757 "mode": self.storage_budget_mode.as_str(),
758 "max_bytes": self.storage_budget_max_bytes,
759 "horizon_years": self.storage_budget_horizon_years,
760 "target_fraction": self.storage_budget_target_fraction,
761 "current_size_bytes": current_size_bytes,
762 "projected_size_bytes": projected_size_bytes,
763 "over_budget": over_budget,
764 "rollup_count": self.storage_budget_rollup_count,
765 },
766 "auto_capture": {
767 "mode": self.auto_capture_mode.as_str(),
768 "redact": self.auto_capture_redact,
769 "max_chars": self.auto_capture_max_chars,
770 "captured_count": self.auto_capture_count
771 },
772 "graph": {
773 "nodes": self.graph.node_count(),
774 "edges": self.graph.edge_count(),
775 "tiers": {
776 "hot": hot,
777 "warm": warm,
778 "cold": cold,
779 },
780 },
781 });
782 let bytes = serde_json::to_vec_pretty(&payload).map_err(|e| {
783 McpError::AgenticMemory(format!("Failed to encode health ledger payload: {e}"))
784 })?;
785 std::fs::write(&tmp, bytes).map_err(McpError::Io)?;
786 std::fs::rename(&tmp, &path).map_err(McpError::Io)?;
787 self.last_health_ledger_emit = Instant::now();
788 Ok(())
789 }
790}
791
792impl Drop for SessionManager {
793 fn drop(&mut self) {
794 if self.dirty {
795 if let Err(e) = self.save() {
796 tracing::error!("Failed to save on drop: {e}");
797 }
798 }
799 if let Err(e) = self.maybe_auto_backup() {
800 tracing::error!("Failed auto-backup on drop: {e}");
801 }
802 }
803}
804
805impl SessionManager {
806 fn auto_archive_completed_sessions(&mut self) -> McpResult<usize> {
807 self.auto_archive_completed_sessions_with_min(self.archive_min_session_nodes)
808 }
809
810 fn auto_archive_completed_sessions_with_min(
811 &mut self,
812 min_session_nodes: usize,
813 ) -> McpResult<usize> {
814 let mut session_ids = self.graph.session_index().session_ids();
815 session_ids.sort_unstable();
816
817 let mut archived = 0usize;
818 for session_id in session_ids {
819 if session_id >= self.current_session {
820 continue;
821 }
822
823 let node_ids = self.graph.session_index().get_session(session_id).to_vec();
824 if node_ids.is_empty() {
825 continue;
826 }
827
828 let mut has_episode = false;
829 let mut event_nodes = 0usize;
830 let mut hot = 0usize;
831 let mut warm = 0usize;
832 let mut cold = 0usize;
833
834 for node_id in &node_ids {
835 if let Some(node) = self.graph.get_node(*node_id) {
836 if node.event_type == EventType::Episode {
837 has_episode = true;
838 continue;
839 }
840 event_nodes += 1;
841 if node.decay_score >= self.hot_min_decay {
842 hot += 1;
843 } else if node.decay_score >= self.warm_min_decay {
844 warm += 1;
845 } else {
846 cold += 1;
847 }
848 }
849 }
850
851 if has_episode || event_nodes < min_session_nodes {
852 continue;
853 }
854
855 let summary = format!(
856 "Auto-archive session {}: {} events ({} hot / {} warm / {} cold)",
857 session_id, event_nodes, hot, warm, cold
858 );
859 self.write_engine
860 .compress_session(&mut self.graph, session_id, &summary)
861 .map_err(|e| {
862 McpError::AgenticMemory(format!(
863 "Auto-archive failed for session {session_id}: {e}"
864 ))
865 })?;
866 archived = archived.saturating_add(1);
867 }
868
869 Ok(archived)
870 }
871
872 fn maybe_enforce_storage_budget(&mut self) -> McpResult<()> {
873 if self.storage_budget_mode == StorageBudgetMode::Off {
874 return Ok(());
875 }
876
877 let current_size = self.current_file_size_bytes();
878 if current_size == 0 {
879 return Ok(());
880 }
881 let projected = self.projected_file_size_bytes(current_size);
882 let over_current = current_size > self.storage_budget_max_bytes;
883 let over_projected = projected
884 .map(|v| v > self.storage_budget_max_bytes)
885 .unwrap_or(false);
886
887 if !over_current && !over_projected {
888 return Ok(());
889 }
890
891 if self.storage_budget_mode == StorageBudgetMode::Warn {
892 tracing::warn!(
893 "Storage budget warning: current={} projected={:?} budget={} (mode=warn)",
894 current_size,
895 projected,
896 self.storage_budget_max_bytes
897 );
898 return Ok(());
899 }
900
901 let target_bytes = ((self.storage_budget_max_bytes as f64
902 * self.storage_budget_target_fraction as f64)
903 .round() as u64)
904 .max(1);
905 let mut rollup_count = 0usize;
906 let mut threshold = self.archive_min_session_nodes.saturating_div(2).max(1);
907
908 for _ in 0..3 {
909 let archived = self.auto_archive_completed_sessions_with_min(threshold)?;
910 if archived == 0 {
911 if threshold > 1 {
912 threshold = 1;
913 continue;
914 }
915 break;
916 }
917 rollup_count += archived;
918 self.dirty = true;
919 self.save()?;
920 let new_size = self.current_file_size_bytes();
921 if new_size <= target_bytes {
922 break;
923 }
924 threshold = 1;
925 }
926
927 if rollup_count > 0 {
928 self.storage_budget_rollup_count = self
929 .storage_budget_rollup_count
930 .saturating_add(rollup_count as u64);
931 tracing::info!(
932 "Storage budget rollup applied: archived_sessions={} budget={} target={} current={}",
933 rollup_count,
934 self.storage_budget_max_bytes,
935 target_bytes,
936 self.current_file_size_bytes()
937 );
938 } else {
939 tracing::warn!(
940 "Storage budget exceeded but no completed sessions eligible for rollup (current={} projected={:?} budget={})",
941 current_size,
942 projected,
943 self.storage_budget_max_bytes
944 );
945 }
946
947 Ok(())
948 }
949
950 fn current_file_size_bytes(&self) -> u64 {
951 std::fs::metadata(&self.file_path)
952 .map(|m| m.len())
953 .unwrap_or(0)
954 }
955
956 fn persist_auto_capture(
957 &mut self,
958 event_type: EventType,
959 raw_text: &str,
960 confidence: f32,
961 ) -> McpResult<Option<u64>> {
962 let mut text = raw_text.trim().to_string();
963 if text.is_empty() {
964 return Ok(None);
965 }
966
967 if self.auto_capture_redact {
968 text = redact_sensitive_tokens(&text);
969 }
970
971 if text.len() > self.auto_capture_max_chars {
972 text.truncate(self.auto_capture_max_chars);
973 text.push_str(" …[truncated]");
974 }
975
976 let (node_id, _) = self.add_event(event_type, &text, confidence, vec![])?;
977 self.auto_capture_count = self.auto_capture_count.saturating_add(1);
978 Ok(Some(node_id))
979 }
980
981 fn projected_file_size_bytes(&self, current_size: u64) -> Option<u64> {
982 if current_size == 0 || self.graph.node_count() < 2 {
983 return None;
984 }
985 let mut min_ts = u64::MAX;
986 let mut max_ts = 0u64;
987 for node in self.graph.nodes() {
988 min_ts = min_ts.min(node.created_at);
989 max_ts = max_ts.max(node.created_at);
990 }
991 if min_ts == u64::MAX || max_ts <= min_ts {
992 return None;
993 }
994
995 let span_secs_raw = (max_ts - min_ts) / 1_000_000;
996 let span_secs = span_secs_raw.max(7 * 24 * 3600) as f64;
998 let per_sec = current_size as f64 / span_secs;
999 let horizon_secs = (self.storage_budget_horizon_years as f64) * 365.25 * 24.0 * 3600.0;
1000 let projected = (per_sec * horizon_secs).round();
1001 Some(projected.max(0.0).min(u64::MAX as f64) as u64)
1002 }
1003
1004 fn tier_counts(&self) -> (usize, usize, usize) {
1005 let mut hot = 0usize;
1006 let mut warm = 0usize;
1007 let mut cold = 0usize;
1008
1009 for node in self.graph.nodes() {
1010 if node.event_type == EventType::Episode {
1011 continue;
1012 }
1013 if node.decay_score >= self.hot_min_decay {
1014 hot += 1;
1015 } else if node.decay_score >= self.warm_min_decay {
1016 warm += 1;
1017 } else {
1018 cold += 1;
1019 }
1020 }
1021
1022 (hot, warm, cold)
1023 }
1024
1025 fn create_migration_checkpoint(&self, from_version: u32) -> McpResult<Option<PathBuf>> {
1026 if !self.file_path.exists() {
1027 return Ok(None);
1028 }
1029
1030 let migration_dir = resolve_migration_dir(&self.file_path);
1031 std::fs::create_dir_all(&migration_dir).map_err(McpError::Io)?;
1032
1033 let ts = chrono::Utc::now().format("%Y%m%d%H%M%S");
1034 let stem = self
1035 .file_path
1036 .file_stem()
1037 .and_then(OsStr::to_str)
1038 .unwrap_or("brain");
1039 let checkpoint = migration_dir.join(format!("{stem}.v{from_version}.{ts}.amem.checkpoint"));
1040 std::fs::copy(&self.file_path, &checkpoint).map_err(McpError::Io)?;
1041 Ok(Some(checkpoint))
1042 }
1043
1044 fn next_backup_path(&self) -> PathBuf {
1045 let ts = chrono::Utc::now().format("%Y%m%d%H%M%S");
1046 let stem = self
1047 .file_path
1048 .file_stem()
1049 .and_then(OsStr::to_str)
1050 .unwrap_or("brain");
1051 self.backups_dir.join(format!("{stem}.{ts}.amem.bak"))
1052 }
1053
1054 fn prune_old_backups(&self) -> McpResult<()> {
1055 let mut entries = std::fs::read_dir(&self.backups_dir)
1056 .map_err(McpError::Io)?
1057 .filter_map(Result::ok)
1058 .filter(|entry| {
1059 entry
1060 .file_name()
1061 .to_str()
1062 .map(|name| name.ends_with(".amem.bak"))
1063 .unwrap_or(false)
1064 })
1065 .collect::<Vec<_>>();
1066
1067 if entries.len() <= self.backup_retention {
1068 return Ok(());
1069 }
1070
1071 entries.sort_by_key(|entry| {
1072 entry
1073 .metadata()
1074 .and_then(|m| m.modified())
1075 .ok()
1076 .unwrap_or(std::time::SystemTime::UNIX_EPOCH)
1077 });
1078 let to_remove = entries.len().saturating_sub(self.backup_retention);
1079 for entry in entries.into_iter().take(to_remove) {
1080 let _ = std::fs::remove_file(entry.path());
1081 }
1082 Ok(())
1083 }
1084}
1085
1086fn resolve_backups_dir(memory_path: &std::path::Path) -> PathBuf {
1087 if let Ok(custom) = std::env::var("AMEM_AUTO_BACKUP_DIR") {
1088 let trimmed = custom.trim();
1089 if !trimmed.is_empty() {
1090 return PathBuf::from(trimmed);
1091 }
1092 }
1093
1094 let parent = memory_path.parent().unwrap_or(std::path::Path::new("."));
1095 parent.join(".amem-backups")
1096}
1097
1098fn resolve_migration_dir(memory_path: &Path) -> PathBuf {
1099 let parent = memory_path.parent().unwrap_or(std::path::Path::new("."));
1100 parent.join(".amem-migrations")
1101}
1102
1103fn read_storage_version(path: &Path) -> Option<u32> {
1104 let mut file = std::fs::File::open(path).ok()?;
1105 let mut header = [0u8; 8];
1106 file.read_exact(&mut header).ok()?;
1107 if &header[0..4] != b"AMEM" {
1108 return None;
1109 }
1110 Some(u32::from_le_bytes([
1111 header[4], header[5], header[6], header[7],
1112 ]))
1113}
1114
1115fn read_env_u64(name: &str, default_value: u64) -> u64 {
1116 std::env::var(name)
1117 .ok()
1118 .and_then(|v| v.parse::<u64>().ok())
1119 .unwrap_or(default_value)
1120}
1121
1122fn read_env_u32(name: &str, default_value: u32) -> u32 {
1123 std::env::var(name)
1124 .ok()
1125 .and_then(|v| v.parse::<u32>().ok())
1126 .unwrap_or(default_value)
1127}
1128
1129fn read_env_usize(name: &str, default_value: usize) -> usize {
1130 std::env::var(name)
1131 .ok()
1132 .and_then(|v| v.parse::<usize>().ok())
1133 .unwrap_or(default_value)
1134}
1135
1136fn read_env_f32(name: &str, default_value: f32) -> f32 {
1137 std::env::var(name)
1138 .ok()
1139 .and_then(|v| v.parse::<f32>().ok())
1140 .unwrap_or(default_value)
1141}
1142
1143fn read_env_bool(name: &str, default_value: bool) -> bool {
1144 std::env::var(name)
1145 .ok()
1146 .map(|v| {
1147 matches!(
1148 v.trim().to_ascii_lowercase().as_str(),
1149 "1" | "true" | "yes" | "on"
1150 )
1151 })
1152 .unwrap_or(default_value)
1153}
1154
1155fn read_env_string(name: &str) -> Option<String> {
1156 std::env::var(name).ok().map(|v| v.trim().to_string())
1157}
1158
1159fn resolve_health_ledger_dir() -> PathBuf {
1160 if let Some(custom) = read_env_string("AMEM_HEALTH_LEDGER_DIR") {
1161 if !custom.is_empty() {
1162 return PathBuf::from(custom);
1163 }
1164 }
1165 if let Some(custom) = read_env_string("AGENTRA_HEALTH_LEDGER_DIR") {
1166 if !custom.is_empty() {
1167 return PathBuf::from(custom);
1168 }
1169 }
1170
1171 let home = std::env::var("HOME")
1172 .ok()
1173 .map(PathBuf::from)
1174 .unwrap_or_else(|| PathBuf::from("."));
1175 home.join(".agentra").join("health-ledger")
1176}
1177
1178fn extract_prompt_capture_text(
1179 prompt_name: &str,
1180 arguments: Option<&Value>,
1181) -> McpResult<Option<String>> {
1182 let args = arguments.unwrap_or(&Value::Null);
1183 let fields = collect_text_fields_by_keys(
1184 args,
1185 &[
1186 "information",
1187 "context",
1188 "topic",
1189 "old_belief",
1190 "new_information",
1191 "reason",
1192 "summary",
1193 "instruction",
1194 "prompt",
1195 "query",
1196 ],
1197 8,
1198 );
1199 if fields.is_empty() {
1200 return Ok(None);
1201 }
1202 let joined = fields.join(" | ");
1203 Ok(Some(format!(
1204 "[auto-capture][prompt] template={prompt_name} input={joined}"
1205 )))
1206}
1207
1208fn extract_safe_tool_capture_text(
1209 tool_name: &str,
1210 arguments: Option<&Value>,
1211) -> McpResult<Option<String>> {
1212 let args = arguments.unwrap_or(&Value::Null);
1213 let keys = ["feedback", "summary", "note"];
1214 if tool_name != "session_end" {
1215 let explicit_feedback = collect_text_fields_by_keys(args, &["feedback", "note"], 4);
1217 if explicit_feedback.is_empty() {
1218 return Ok(None);
1219 }
1220 }
1221 let fields = collect_text_fields_by_keys(args, &keys, 6);
1222 if fields.is_empty() {
1223 return Ok(None);
1224 }
1225 Ok(Some(format!(
1226 "[auto-capture][feedback] tool={tool_name} context={}",
1227 fields.join(" | ")
1228 )))
1229}
1230
1231fn extract_full_tool_capture_text(
1232 tool_name: &str,
1233 arguments: Option<&Value>,
1234) -> McpResult<Option<String>> {
1235 if tool_name == "memory_add" {
1236 return Ok(None);
1237 }
1238 let args = arguments.unwrap_or(&Value::Null);
1239 let preferred = collect_text_fields_by_keys(
1240 args,
1241 &[
1242 "query",
1243 "content",
1244 "prompt",
1245 "new_content",
1246 "reason",
1247 "summary",
1248 "topic",
1249 "instruction",
1250 "information",
1251 "context",
1252 "feedback",
1253 ],
1254 10,
1255 );
1256
1257 let fields = if preferred.is_empty() {
1258 collect_all_string_like_fields(args, 8)
1259 } else {
1260 preferred
1261 };
1262
1263 if fields.is_empty() {
1264 return Ok(None);
1265 }
1266 Ok(Some(format!(
1267 "[auto-capture][tool] tool={tool_name} input={}",
1268 fields.join(" | ")
1269 )))
1270}
1271
1272fn collect_text_fields_by_keys(value: &Value, keys: &[&str], limit: usize) -> Vec<String> {
1273 let mut out = Vec::new();
1274 let mut seen = std::collections::BTreeSet::<String>::new();
1275
1276 fn walk(
1277 value: &Value,
1278 path: String,
1279 keys: &[&str],
1280 out: &mut Vec<String>,
1281 seen: &mut std::collections::BTreeSet<String>,
1282 limit: usize,
1283 ) {
1284 if out.len() >= limit {
1285 return;
1286 }
1287 match value {
1288 Value::Object(map) => {
1289 for (k, v) in map {
1290 if out.len() >= limit {
1291 break;
1292 }
1293 let next = if path.is_empty() {
1294 k.to_string()
1295 } else {
1296 format!("{path}.{k}")
1297 };
1298 let key_match = keys
1299 .iter()
1300 .any(|needle| k.eq_ignore_ascii_case(needle) || next.ends_with(needle));
1301 if key_match {
1302 if let Some(s) = value_to_compact_string(v) {
1303 let entry = format!("{next}={s}");
1304 if seen.insert(entry.clone()) {
1305 out.push(entry);
1306 }
1307 }
1308 }
1309 walk(v, next, keys, out, seen, limit);
1310 }
1311 }
1312 Value::Array(items) => {
1313 for (idx, item) in items.iter().enumerate() {
1314 if out.len() >= limit {
1315 break;
1316 }
1317 let next = format!("{path}[{idx}]");
1318 walk(item, next, keys, out, seen, limit);
1319 }
1320 }
1321 _ => {}
1322 }
1323 }
1324
1325 walk(value, String::new(), keys, &mut out, &mut seen, limit);
1326 out
1327}
1328
1329fn collect_all_string_like_fields(value: &Value, limit: usize) -> Vec<String> {
1330 let mut out = Vec::new();
1331 fn walk(value: &Value, path: String, out: &mut Vec<String>, limit: usize) {
1332 if out.len() >= limit {
1333 return;
1334 }
1335 match value {
1336 Value::Object(map) => {
1337 for (k, v) in map {
1338 if out.len() >= limit {
1339 break;
1340 }
1341 let next = if path.is_empty() {
1342 k.to_string()
1343 } else {
1344 format!("{path}.{k}")
1345 };
1346 walk(v, next, out, limit);
1347 }
1348 }
1349 Value::Array(items) => {
1350 for (idx, item) in items.iter().enumerate() {
1351 if out.len() >= limit {
1352 break;
1353 }
1354 walk(item, format!("{path}[{idx}]"), out, limit);
1355 }
1356 }
1357 _ => {
1358 if let Some(s) = value_to_compact_string(value) {
1359 out.push(format!("{path}={s}"));
1360 }
1361 }
1362 }
1363 }
1364 walk(value, String::new(), &mut out, limit);
1365 out
1366}
1367
1368fn value_to_compact_string(value: &Value) -> Option<String> {
1369 match value {
1370 Value::String(s) => {
1371 let trimmed = s.trim();
1372 if trimmed.is_empty() {
1373 None
1374 } else {
1375 Some(trimmed.to_string())
1376 }
1377 }
1378 Value::Number(n) => Some(n.to_string()),
1379 Value::Bool(b) => Some(b.to_string()),
1380 Value::Null => None,
1381 Value::Array(arr) => {
1382 if arr.is_empty() {
1383 None
1384 } else {
1385 Some(format!("<array:{}>", arr.len()))
1386 }
1387 }
1388 Value::Object(map) => {
1389 if map.is_empty() {
1390 None
1391 } else {
1392 Some(format!("<object:{}>", map.len()))
1393 }
1394 }
1395 }
1396}
1397
1398fn redact_sensitive_tokens(text: &str) -> String {
1399 text.split_whitespace()
1400 .map(redact_token)
1401 .collect::<Vec<_>>()
1402 .join(" ")
1403}
1404
1405fn redact_token(token: &str) -> String {
1406 let trimmed = token.trim_matches(|c: char| c == '"' || c == '\'' || c == ',' || c == ';');
1407 let lower = trimmed.to_ascii_lowercase();
1408 if trimmed.starts_with("/Users/")
1409 || trimmed.starts_with("C:\\Users\\")
1410 || trimmed.contains("/Users/")
1411 || trimmed.contains("C:\\Users\\")
1412 {
1413 return "[REDACTED_PATH]".to_string();
1414 }
1415 if trimmed.contains('@') && trimmed.contains('.') {
1416 return "[REDACTED_EMAIL]".to_string();
1417 }
1418 if lower.starts_with("sk-")
1419 || lower.contains("api_key")
1420 || lower.contains("access_token")
1421 || lower.contains("bearer")
1422 || lower.contains("authorization")
1423 {
1424 return "[REDACTED_SECRET]".to_string();
1425 }
1426 if looks_like_long_secret(trimmed) {
1427 return "[REDACTED_SECRET]".to_string();
1428 }
1429 token.to_string()
1430}
1431
1432fn looks_like_long_secret(token: &str) -> bool {
1433 if token.len() < 24 {
1434 return false;
1435 }
1436 token
1437 .chars()
1438 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')
1439}
1440
1441#[cfg(test)]
1442mod tests {
1443 use super::*;
1444 use serde_json::json;
1445
1446 #[test]
1447 fn budget_projection_available_with_timeline() {
1448 let dir = tempfile::tempdir().expect("tempdir");
1449 let brain = dir.path().join("projection.amem");
1450 let mut manager = SessionManager::open(brain.to_str().expect("path")).expect("open");
1451
1452 let (id_a, _) = manager
1453 .add_event(EventType::Fact, "old fact", 0.9, vec![])
1454 .expect("add fact");
1455 let (_id_b, _) = manager
1456 .add_event(EventType::Fact, "new fact", 0.9, vec![])
1457 .expect("add fact");
1458
1459 {
1460 let graph = manager.graph_mut();
1461 let old = graph.get_node_mut(id_a).expect("node");
1462 old.created_at = old.created_at.saturating_sub(15 * 24 * 3600 * 1_000_000);
1463 }
1464 manager.save().expect("save");
1465 let size = manager.current_file_size_bytes();
1466 let projected = manager.projected_file_size_bytes(size);
1467 assert!(size > 0);
1468 assert!(projected.is_some());
1469 }
1470
1471 #[test]
1472 fn budget_auto_rollup_archives_completed_session() {
1473 let dir = tempfile::tempdir().expect("tempdir");
1474 let brain = dir.path().join("rollup.amem");
1475 let mut manager = SessionManager::open(brain.to_str().expect("path")).expect("open");
1476
1477 let _ = manager
1479 .add_event(EventType::Fact, "alpha", 0.8, vec![])
1480 .expect("add");
1481 let _ = manager
1482 .add_event(EventType::Decision, "beta", 0.9, vec![])
1483 .expect("add");
1484 manager.start_session(Some(2)).expect("session");
1485 manager.save().expect("save");
1486
1487 manager.storage_budget_mode = StorageBudgetMode::AutoRollup;
1489 manager.storage_budget_max_bytes = 1;
1490 manager.storage_budget_target_fraction = 0.5;
1491
1492 manager
1493 .maybe_enforce_storage_budget()
1494 .expect("enforce budget");
1495
1496 let episode_count = manager
1497 .graph()
1498 .nodes()
1499 .iter()
1500 .filter(|n| n.event_type == EventType::Episode)
1501 .count();
1502 assert!(episode_count >= 1);
1503 assert!(manager.storage_budget_rollup_count >= 1);
1504 }
1505
1506 #[test]
1507 fn auto_capture_off_noop() {
1508 let dir = tempfile::tempdir().expect("tempdir");
1509 let brain = dir.path().join("capture-off.amem");
1510 let mut manager = SessionManager::open(brain.to_str().expect("path")).expect("open");
1511 manager.auto_capture_mode = AutoCaptureMode::Off;
1512
1513 let captured = manager
1514 .capture_prompt_request(
1515 "remember",
1516 Some(&json!({"information":"hello world","context":"ctx"})),
1517 )
1518 .expect("capture");
1519 assert!(captured.is_none());
1520 assert_eq!(manager.graph().node_count(), 0);
1521 }
1522
1523 #[test]
1524 fn auto_capture_full_records_and_redacts() {
1525 let dir = tempfile::tempdir().expect("tempdir");
1526 let brain = dir.path().join("capture-full.amem");
1527 let mut manager = SessionManager::open(brain.to_str().expect("path")).expect("open");
1528 manager.auto_capture_mode = AutoCaptureMode::Full;
1529 manager.auto_capture_redact = true;
1530
1531 manager
1532 .capture_tool_call(
1533 "memory_query",
1534 Some(&json!({
1535 "query":"Find anything about token sk-THISISALONGSECRET123456",
1536 "context":"/Users/omoshola/Documents/private.txt",
1537 "reason":"email me at test@example.com"
1538 })),
1539 )
1540 .expect("capture");
1541
1542 assert!(manager.graph().node_count() >= 1);
1543 let latest = manager
1544 .graph()
1545 .nodes()
1546 .iter()
1547 .max_by_key(|n| n.id)
1548 .expect("node");
1549 assert!(latest.content.contains("[auto-capture][tool]"));
1550 assert!(latest.content.contains("[REDACTED_SECRET]"));
1551 assert!(latest.content.contains("[REDACTED_PATH]"));
1552 assert!(latest.content.contains("[REDACTED_EMAIL]"));
1553 }
1554}