Skip to main content

agentic_memory_mcp/session/
manager.rs

1//! Graph lifecycle management, file I/O, and session tracking.
2
3use 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
16/// Default auto-save interval.
17const DEFAULT_AUTO_SAVE_SECS: u64 = 30;
18/// Default backup interval.
19const DEFAULT_BACKUP_INTERVAL_SECS: u64 = 900;
20/// Default number of backups to retain per brain file.
21const DEFAULT_BACKUP_RETENTION: usize = 24;
22/// Default maintenance sleep-cycle interval.
23const DEFAULT_SLEEP_CYCLE_SECS: u64 = 1800;
24/// Minimum completed-session size before auto-archive.
25const DEFAULT_ARCHIVE_MIN_SESSION_NODES: usize = 25;
26/// Default hot-tier threshold (decay score).
27const DEFAULT_HOT_MIN_DECAY: f32 = 0.7;
28/// Default warm-tier threshold (decay score).
29const DEFAULT_WARM_MIN_DECAY: f32 = 0.3;
30/// Default sustained mutation rate threshold before throttling heavy maintenance.
31const DEFAULT_SLA_MAX_MUTATIONS_PER_MIN: u32 = 240;
32/// Default interval for writing health-ledger snapshots.
33const DEFAULT_HEALTH_LEDGER_EMIT_SECS: u64 = 30;
34/// Default long-horizon storage budget target (2 GiB over 20 years).
35const DEFAULT_STORAGE_BUDGET_BYTES: u64 = 2 * 1024 * 1024 * 1024;
36/// Default storage budget projection horizon.
37const DEFAULT_STORAGE_BUDGET_HORIZON_YEARS: u32 = 20;
38/// Default maximum chars persisted for one auto-captured prompt/feedback item.
39const DEFAULT_AUTO_CAPTURE_MAX_CHARS: usize = 2048;
40/// Current `.amem` storage version used by this server.
41const 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    /// Capture prompt-focused events and feedback context.
67    Safe,
68    /// Capture broader tool input text (except explicit memory_add payload duplication).
69    Full,
70    /// Disable automatic capture.
71    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
201/// Manages the memory graph lifecycle, file I/O, and session state.
202pub 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    /// Open or create a memory file at the given path.
245    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            // Ensure parent directory exists
266            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        // Determine the next session ID from existing sessions
278        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    /// Get an immutable reference to the graph.
423    pub fn graph(&self) -> &MemoryGraph {
424        &self.graph
425    }
426
427    /// Get a mutable reference to the graph and mark as dirty.
428    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    /// Get the query engine.
436    pub fn query_engine(&self) -> &QueryEngine {
437        &self.query_engine
438    }
439
440    /// Get the write engine.
441    pub fn write_engine(&self) -> &WriteEngine {
442        &self.write_engine
443    }
444
445    /// Current session ID.
446    pub fn current_session_id(&self) -> u32 {
447        self.current_session
448    }
449
450    /// Start a new session, optionally with an explicit ID.
451    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    /// End a session and optionally create an episode summary.
464    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    /// Save the graph to file.
481    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    /// Check if auto-save is needed and save if so.
499    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    /// Runs autonomous maintenance: sleep-cycle, auto-save, and periodic backup.
507    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    /// Run a periodic sleep-cycle: decay refresh + tier balancing + auto-archive.
529    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    /// Periodic backup of persisted state with retention pruning.
563    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    /// Mark the graph as dirty (needs saving).
585    pub fn mark_dirty(&mut self) {
586        self.dirty = true;
587        self.last_activity = Instant::now();
588        self.record_mutation();
589    }
590
591    /// Get the file path.
592    pub fn file_path(&self) -> &PathBuf {
593        &self.file_path
594    }
595
596    /// Background maintenance loop interval.
597    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    /// Capture a prompt template invocation (`prompts/get`) into memory.
604    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    /// Capture a tool call input context into memory based on capture mode.
619    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    /// Add a cognitive event to the graph.
640    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        // First, add the node to get its assigned ID
653        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        // Then add edges with the correct source_id
663        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    /// Correct a previous belief.
681    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        // Clamp to at least one week to avoid unstable extrapolation on tiny windows.
997        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        // Keep safe mode low-noise and non-invasive: only capture explicit feedback fields.
1216        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        // Build session 1 with enough content, then move to session 2 so session 1 is completed.
1478        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        // Force tiny budget to trigger rollup.
1488        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}