1use std::collections::{HashMap, HashSet};
27
28use mentedb_core::error::MenteResult;
29use mentedb_core::memory::{AttributeValue, MemoryNode, MemoryType};
30use mentedb_core::types::{AgentId, MemoryId, UserId};
31use mentedb_storage::PageId;
32
33use crate::MenteDb;
34
35pub const ATTR_INJECTION_SHOWN: &str = "injection_shown";
37pub const ATTR_INJECTION_USED: &str = "injection_used";
40
41#[derive(Debug, Clone)]
43pub struct InjectionConfig {
44 pub candidate_pool: usize,
46 pub cluster_dominant_fraction: f32,
55 pub cluster_fill_max: usize,
57 pub cluster_max_span: usize,
61 pub use_lexical_leg: bool,
66 pub own_agent_boost: f32,
72 pub escalation_boost: f32,
79 pub escalation_min_relevance: f32,
84 pub mmr_lambda: f32,
86 pub knee_gap_ratio: f32,
88 pub demotion_shown_min: i64,
90 pub demotion_factor: f32,
92 pub demotion_drop_shown: i64,
96 pub tail_min_keep: usize,
100 pub tail_floor_fraction: f32,
106 pub cross_project_episodic: bool,
111 pub cross_project_floor_fraction: f32,
116 pub used_similarity: f32,
118 pub use_reinforcement: f32,
120 pub graph_expansion_max: usize,
125 pub graph_expansion_decay: f32,
129 pub mode_trigger_gate: f32,
136 pub mode_rules_max: usize,
138 pub mode_probe_load_budget: usize,
145}
146
147impl Default for InjectionConfig {
148 fn default() -> Self {
149 Self {
150 candidate_pool: 48,
151 cluster_dominant_fraction: 0.0,
152 cluster_fill_max: 4,
153 cluster_max_span: 12,
154 use_lexical_leg: true,
155 own_agent_boost: 1.25,
156 escalation_boost: 2.5,
157 escalation_min_relevance: 0.5,
158 mmr_lambda: 0.7,
159 knee_gap_ratio: 2.0,
160 demotion_shown_min: 5,
161 demotion_factor: 0.5,
162 demotion_drop_shown: 12,
163 tail_min_keep: 3,
164 tail_floor_fraction: 0.55,
165 cross_project_episodic: false,
166 cross_project_floor_fraction: 0.8,
167 used_similarity: 0.6,
168 use_reinforcement: 0.05,
169 graph_expansion_max: 0,
170 graph_expansion_decay: 0.5,
171 mode_trigger_gate: 0.40,
172 mode_rules_max: 4,
173 mode_probe_load_budget: 256,
174 }
175 }
176}
177
178pub struct InjectionQuery<'a> {
180 pub embedding: &'a [f32],
182 pub query_text: Option<&'a str>,
184 pub session_id: Option<&'a str>,
187 pub exclude_ids: &'a [MemoryId],
189 pub max_items: usize,
191 pub max_episodic: usize,
193 pub agent_id: Option<AgentId>,
196 pub user_id: Option<UserId>,
200 pub current_project: Option<&'a str>,
205}
206
207#[derive(Debug, Clone, Copy, PartialEq, Eq)]
209pub enum SelectionReason {
210 Pinned,
212 Relevant,
214 ModeRule,
217}
218
219pub struct InjectionCandidate {
221 pub node: MemoryNode,
222 pub score: f32,
225 pub reason: SelectionReason,
226}
227
228fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
229 if a.len() != b.len() || a.is_empty() {
230 return 0.0;
231 }
232 let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
233 let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
234 let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
235 if na == 0.0 || nb == 0.0 {
236 return 0.0;
237 }
238 dot / (na * nb)
239}
240
241fn attr_count(node: &MemoryNode, key: &str) -> i64 {
242 match node.attributes.get(key) {
243 Some(AttributeValue::Integer(n)) => *n,
244 _ => 0,
245 }
246}
247
248fn bump_attr(node: &mut MemoryNode, key: &str) {
249 let next = attr_count(node, key) + 1;
250 node.attributes
251 .insert(key.to_string(), AttributeValue::Integer(next));
252}
253
254fn has_tag(node: &MemoryNode, tag: &str) -> bool {
255 node.tags.iter().any(|t| t == tag)
256}
257
258fn now_us() -> u64 {
259 std::time::SystemTime::now()
260 .duration_since(std::time::UNIX_EPOCH)
261 .unwrap_or_default()
262 .as_micros() as u64
263}
264
265fn knee_cutoff(scores: &[f32], gap_ratio: f32) -> usize {
268 if scores.len() <= 1 {
269 return scores.len();
270 }
271 let mut best_ratio = 0.0f32;
272 let mut cut = scores.len();
273 for i in 0..scores.len() - 1 {
274 let hi = scores[i];
275 let lo = scores[i + 1].max(f32::EPSILON);
276 let ratio = hi / lo;
277 if ratio > best_ratio {
278 best_ratio = ratio;
279 cut = i + 1;
280 }
281 }
282 if best_ratio >= gap_ratio {
283 cut
284 } else {
285 scores.len()
286 }
287}
288
289fn project_of(node: &MemoryNode) -> Option<&str> {
291 node.tags
292 .iter()
293 .find_map(|t| t.strip_prefix("scope:project:"))
294}
295
296fn cross_project_allowed(
302 node: &MemoryNode,
303 cos: f32,
304 best_cos: f32,
305 current_project: Option<&str>,
306 cfg: &InjectionConfig,
307) -> bool {
308 let Some(current) = current_project else {
309 return true;
310 };
311 let Some(mem_project) = project_of(node) else {
312 return true;
313 };
314 if mem_project == current {
315 return true;
316 }
317 if node.memory_type == MemoryType::Episodic {
318 return cfg.cross_project_episodic;
319 }
320 cos >= best_cos * cfg.cross_project_floor_fraction
321}
322
323fn passes_tail_gate(idx: usize, cos: f32, best_cos: f32, cfg: &InjectionConfig) -> bool {
330 if idx < cfg.tail_min_keep {
331 return true;
332 }
333 cos >= best_cos * cfg.tail_floor_fraction
334}
335
336impl MenteDb {
337 pub fn recall_for_injection(
340 &self,
341 query: &InjectionQuery<'_>,
342 ) -> MenteResult<Vec<InjectionCandidate>> {
343 let cfg = self.cognitive_config.injection_config.clone();
344 let excluded: HashSet<MemoryId> = query.exclude_ids.iter().copied().collect();
345 let session_tag = query.session_id.map(|s| format!("session:{s}"));
346
347 let t_start = std::time::Instant::now();
350 let hits = self
352 .recall_hybrid_scoped_at_mode(
353 query.embedding,
354 query.query_text.filter(|_| cfg.use_lexical_leg),
355 cfg.candidate_pool,
356 now_us(),
357 None,
358 false,
359 None,
360 query.agent_id,
361 query.user_id,
362 query.current_project,
363 )
364 .unwrap_or_default();
365
366 let t_hybrid = t_start.elapsed();
367 let hybrid_hits = hits.len();
368 let mut candidates: Vec<(MemoryNode, f32, f32)> = Vec::new();
369 for (id, score) in hits {
370 if excluded.contains(&id) {
371 continue;
372 }
373 let Ok(node) = self.get_memory(id) else {
374 continue;
375 };
376 if has_tag(&node, "action")
377 || has_tag(&node, "scope:always")
378 || has_tag(&node, "ghost-memory")
379 || node.tags.iter().any(|t| t.starts_with("mode-exemplar:"))
380 {
381 continue;
386 }
387 if let Some(ref st) = session_tag
388 && has_tag(&node, st)
389 {
390 continue;
391 }
392 let shown = attr_count(&node, ATTR_INJECTION_SHOWN);
398 let used = attr_count(&node, ATTR_INJECTION_USED);
399 if shown >= cfg.demotion_drop_shown && used == 0 {
400 continue;
401 }
402 let mut adjusted = if shown >= cfg.demotion_shown_min && used == 0 {
403 score * cfg.demotion_factor
404 } else {
405 score
406 };
407 if let Some(agent) = query.agent_id
408 && node.agent_id == agent
409 && !agent.0.is_nil()
410 {
411 adjusted *= cfg.own_agent_boost;
412 }
413 let cos = cosine_similarity(query.embedding, &node.embedding);
414 candidates.push((node, adjusted, cos));
415 }
416
417 let cluster_pool: Vec<(MemoryNode, f32)> = if cfg.cluster_dominant_fraction > 0.0 {
420 candidates
421 .iter()
422 .filter(|(node, _, _)| node.tags.iter().any(|t| t.starts_with("section:")))
423 .map(|(node, _, cos)| (node.clone(), *cos))
424 .collect()
425 } else {
426 Vec::new()
427 };
428
429 let best_cos = candidates.iter().map(|(_, _, c)| *c).fold(0.0f32, f32::max);
433 candidates.retain(|(node, _, cos)| {
434 cross_project_allowed(node, *cos, best_cos, query.current_project, &cfg)
435 });
436
437 candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
438 let scores: Vec<f32> = candidates.iter().map(|(_, s, _)| *s).collect();
439 candidates.truncate(knee_cutoff(&scores, cfg.knee_gap_ratio));
440
441 let mut idx = 0usize;
445 candidates.retain(|(_, _, cos)| {
446 let keep = passes_tail_gate(idx, *cos, best_cos, &cfg);
447 idx += 1;
448 keep
449 });
450
451 if cfg.escalation_boost != 1.0
457 && candidates
458 .iter()
459 .any(|(n, _, _)| has_tag(n, "escalated:violation"))
460 {
461 let relevance_floor = best_cos * cfg.escalation_min_relevance;
462 for (node, adjusted, cos) in candidates.iter_mut() {
463 if has_tag(node, "escalated:violation") && *cos >= relevance_floor {
464 *adjusted *= cfg.escalation_boost;
465 }
466 }
467 candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
468 }
469
470 let mut scored: Vec<(MemoryNode, f32)> =
471 candidates.into_iter().map(|(n, s, _)| (n, s)).collect();
472
473 if cfg.graph_expansion_max > 0 && !scored.is_empty() {
484 let now = now_us();
485 let mut present: HashSet<MemoryId> = scored.iter().map(|(n, _)| n.id).collect();
486 present.extend(excluded.iter().copied());
487 let seeds: Vec<(MemoryId, f32)> = scored.iter().map(|(n, s)| (n.id, *s)).collect();
488
489 let g = self.graph.graph();
490 let mut added = 0usize;
491 'seeds: for (seed_id, seed_score) in seeds {
492 let mut edges = g.outgoing(seed_id);
493 edges.extend(g.incoming(seed_id));
494 for (nbr_id, edge) in edges {
495 if added >= cfg.graph_expansion_max {
496 break 'seeds;
497 }
498 if present.contains(&nbr_id) {
499 continue;
500 }
501 let Ok(node) = self.get_memory(nbr_id) else {
502 continue;
503 };
504 if has_tag(&node, "action")
505 || has_tag(&node, "scope:always")
506 || has_tag(&node, "ghost-memory")
507 || !node.is_valid_at(now)
508 || !crate::agent_visible(node.agent_id, query.agent_id)
509 || !crate::user_visible(node.user_id, query.user_id)
510 {
511 continue;
512 }
513 if let Some(ref st) = session_tag
514 && has_tag(&node, st)
515 {
516 continue;
517 }
518 let score = seed_score * edge.weight * cfg.graph_expansion_decay;
519 present.insert(nbr_id);
520 scored.push((node, score));
521 added += 1;
522 }
523 }
524 }
525
526 let dominant_section: Option<String> = if cfg.cluster_dominant_fraction > 0.0 {
530 let head = scored.len().min(query.max_items.max(3));
531 let mut counts: HashMap<String, usize> = HashMap::new();
532 for (node, _) in &scored[..head] {
533 for tag in &node.tags {
534 if let Some(sec) = tag.strip_prefix("section:") {
535 *counts.entry(sec.to_string()).or_insert(0) += 1;
536 }
537 }
538 }
539 let mut span: HashMap<&str, usize> = HashMap::new();
543 for (node, _) in &cluster_pool {
544 for tag in &node.tags {
545 if let Some(sec) = tag.strip_prefix("section:") {
546 *span.entry(sec).or_insert(0) += 1;
547 }
548 }
549 }
550 let need = ((cfg.cluster_dominant_fraction * head as f32).ceil() as usize).max(2);
551 counts
552 .into_iter()
553 .filter(|(sec, n)| {
554 *n >= need
555 && span.get(sec.as_str()).copied().unwrap_or(0) <= cfg.cluster_max_span
556 })
557 .max_by_key(|(_, n)| *n)
558 .map(|(sec, _)| sec)
559 } else {
560 None
561 };
562
563 let top_score = scored
565 .first()
566 .map(|(_, s)| *s)
567 .unwrap_or(1.0)
568 .max(f32::EPSILON);
569 let mut remaining: Vec<(MemoryNode, f32)> = scored;
570 let mut selected: Vec<InjectionCandidate> = Vec::new();
571 let mut episodic_count = 0usize;
572
573 while selected.len() < query.max_items && !remaining.is_empty() {
574 let mut best_idx: Option<usize> = None;
575 let mut best_value = f32::NEG_INFINITY;
576 for (idx, (node, score)) in remaining.iter().enumerate() {
577 if node.memory_type == MemoryType::Episodic && episodic_count >= query.max_episodic
578 {
579 continue;
580 }
581 let redundancy = selected
582 .iter()
583 .map(|s| cosine_similarity(&node.embedding, &s.node.embedding))
584 .fold(0.0f32, f32::max);
585 let value =
586 cfg.mmr_lambda * (score / top_score) - (1.0 - cfg.mmr_lambda) * redundancy;
587 if value > best_value {
588 best_value = value;
589 best_idx = Some(idx);
590 }
591 }
592 let Some(idx) = best_idx else {
593 break;
594 };
595 let (node, score) = remaining.remove(idx);
596 if node.memory_type == MemoryType::Episodic {
597 episodic_count += 1;
598 }
599 selected.push(InjectionCandidate {
600 node,
601 score,
602 reason: SelectionReason::Relevant,
603 });
604 }
605
606 if selected.len() >= 3
614 && let Some(sec) = dominant_section
615 {
616 {
617 let sec_tag = format!("section:{sec}");
618 let mut siblings: Vec<(MemoryNode, f32)> = cluster_pool
619 .into_iter()
620 .filter(|(node, _)| {
621 node.tags.iter().any(|t| t == &sec_tag)
622 && !selected.iter().any(|c| c.node.id == node.id)
623 })
624 .collect();
625 siblings.sort_by(|a, b| b.1.total_cmp(&a.1));
626 siblings.truncate(cfg.cluster_fill_max);
627 let mut evict: Vec<usize> = (1..selected.len())
632 .filter(|i| !selected[*i].node.tags.iter().any(|t| t == &sec_tag))
633 .collect();
634 evict.sort_by(|a, b| {
635 selected[*a]
636 .score
637 .partial_cmp(&selected[*b].score)
638 .unwrap_or(std::cmp::Ordering::Equal)
639 });
640 for (pos, (node, cos)) in evict.into_iter().zip(siblings) {
641 selected[pos] = InjectionCandidate {
642 node,
643 score: cos,
644 reason: SelectionReason::Relevant,
645 };
646 }
647 }
648 }
649
650 let t_selected = t_start.elapsed();
657 let mut pinned_loads = 0usize;
658 let mut result: Vec<InjectionCandidate> = Vec::new();
659 for mid in self.index.bitmap.query_tag("scope:always") {
660 if excluded.contains(&mid) {
661 continue;
662 }
663 pinned_loads += 1;
664 if let Ok(node) = self.get_memory(mid)
665 && crate::agent_visible(node.agent_id, query.agent_id)
666 && crate::user_visible(node.user_id, query.user_id)
667 {
668 result.push(InjectionCandidate {
669 node,
670 score: 0.0,
671 reason: SelectionReason::Pinned,
672 });
673 }
674 }
675 let t_pinned = t_start.elapsed();
681 let mut mode_tags = 0usize;
682 let mut mode_loads = 0usize;
683 if cfg.mode_trigger_gate > 0.0 && !query.embedding.is_empty() {
684 let now = now_us();
685 let mut mode_rules: Vec<(f32, MemoryNode)> = Vec::new();
686 let mut mode_candidates: Vec<(f32, MemoryId, String)> = Vec::new();
687 for tag in self.index.bitmap.tags_with_prefix("trigger:") {
688 mode_tags += 1;
689 let trigger = tag.trim_start_matches("trigger:");
690 let exemplar_tag = format!("mode-exemplar:{trigger}");
695 let mut activation = f32::MIN;
696 let exemplar_ids = self.index.bitmap.query_tag(&exemplar_tag);
697 if exemplar_ids.is_empty() {
698 for mid in self.index.bitmap.query_tag(&tag) {
699 if let Some(probe) = self.mode_probe(mid, &mut mode_loads)
700 && probe.is_valid_at(now)
701 && !probe.embedding.is_empty()
702 && crate::agent_visible(probe.agent_id, query.agent_id)
703 && crate::user_visible(probe.user_id, query.user_id)
704 {
705 let sim = cosine_similarity(query.embedding, &probe.embedding);
706 activation = activation.max(sim);
707 }
708 }
709 } else {
710 for mid in exemplar_ids {
711 if let Some(probe) = self.mode_probe(mid, &mut mode_loads)
712 && probe.is_valid_at(now)
713 && !probe.embedding.is_empty()
714 && crate::agent_visible(probe.agent_id, query.agent_id)
715 && crate::user_visible(probe.user_id, query.user_id)
716 {
717 let sim = cosine_similarity(query.embedding, &probe.embedding);
718 activation = activation.max(sim);
719 }
720 }
721 }
722 if activation < cfg.mode_trigger_gate {
723 continue;
724 }
725 for mid in self.index.bitmap.query_tag(&tag) {
729 if excluded.contains(&mid)
730 || result.iter().any(|c| c.node.id == mid)
731 || selected.iter().any(|c| c.node.id == mid)
732 || mode_candidates.iter().any(|(_, id, _)| *id == mid)
733 {
734 continue;
735 }
736 let Some(probe) = self.mode_probe(mid, &mut mode_loads) else {
737 continue;
738 };
739 if !probe.is_valid_at(now)
740 || !crate::agent_visible(probe.agent_id, query.agent_id)
741 || !crate::user_visible(probe.user_id, query.user_id)
742 {
743 continue;
744 }
745 mode_candidates.push((activation, mid, tag.clone()));
746 }
747 }
748 mode_candidates.sort_by(|a, b| b.0.total_cmp(&a.0));
753 for (activation, mid, tag) in mode_candidates {
754 if mode_rules.len() >= cfg.mode_rules_max {
755 break;
756 }
757 mode_loads += 1;
758 let Ok(node) = self.get_memory(mid) else {
759 continue;
760 };
761 if node.tags.iter().any(|t| t == &tag) {
762 mode_rules.push((activation, node));
763 }
764 }
765 for (sim, node) in mode_rules {
766 result.push(InjectionCandidate {
767 node,
768 score: sim,
769 reason: SelectionReason::ModeRule,
770 });
771 }
772 }
773
774 result.extend(selected);
775 let total = t_start.elapsed();
776 if total >= std::time::Duration::from_millis(self.cognitive_config.slow_recall_log_ms) {
777 tracing::info!(
778 total_ms = total.as_millis() as u64,
779 hybrid_ms = t_hybrid.as_millis() as u64,
780 select_ms = (t_selected - t_hybrid).as_millis() as u64,
781 pinned_ms = (t_pinned - t_selected).as_millis() as u64,
782 mode_ms = (total - t_pinned).as_millis() as u64,
783 hybrid_hits,
784 pinned_loads,
785 mode_tags,
786 mode_loads,
787 "slow injection recall"
788 );
789 }
790 Ok(result)
791 }
792
793 pub fn record_injection_outcome(
798 &self,
799 shown: &[MemoryId],
800 reply_embedding: Option<&[f32]>,
801 ) -> MenteResult<(usize, usize)> {
802 let cfg = self.cognitive_config.injection_config.clone();
803 let mut updated = 0usize;
804 let mut used_total = 0usize;
805 let now = now_us();
806
807 for id in shown {
808 let pid = {
809 let pm = self.page_map.read();
810 match pm.get(id) {
811 Some(p) => *p,
812 None => continue,
813 }
814 };
815 let Ok(mut node) = self.storage.load_memory(pid) else {
816 continue;
817 };
818 bump_attr(&mut node, ATTR_INJECTION_SHOWN);
819
820 node.access_count = node.access_count.saturating_add(1);
828 node.accessed_at = now;
829
830 let used = reply_embedding
831 .map(|re| cosine_similarity(&node.embedding, re) >= cfg.used_similarity)
832 .unwrap_or(false);
833 if used {
834 bump_attr(&mut node, ATTR_INJECTION_USED);
837 node.salience = (node.salience + cfg.use_reinforcement).min(1.0);
838 used_total += 1;
839 }
840
841 self.update_node(pid, &node)?;
844 updated += 1;
845 }
846 Ok((updated, used_total))
847 }
848}
849
850pub(crate) struct ModeProbe {
854 pub(crate) embedding: Vec<f32>,
855 pub(crate) agent_id: AgentId,
856 pub(crate) user_id: UserId,
857 valid_from: Option<u64>,
858 valid_until: Option<u64>,
859 stamp: Option<(PageId, u64)>,
865}
866
867impl ModeProbe {
868 fn from_node(node: &MemoryNode, stamp: Option<(PageId, u64)>) -> Self {
869 Self {
870 embedding: node.embedding.clone(),
871 agent_id: node.agent_id,
872 user_id: node.user_id,
873 valid_from: node.valid_from,
874 valid_until: node.valid_until,
875 stamp,
876 }
877 }
878
879 pub(crate) fn is_valid_at(&self, at: u64) -> bool {
881 let from = self.valid_from.unwrap_or(0);
882 match self.valid_until {
883 Some(until) => at >= from && at < until,
884 None => at >= from,
885 }
886 }
887}
888
889impl MenteDb {
890 pub(crate) fn mode_probe(
895 &self,
896 id: MemoryId,
897 loads: &mut usize,
898 ) -> Option<std::sync::Arc<ModeProbe>> {
899 if let Some(ProbeSlot::Ready(p)) = self.mode_probes.read().get(&id) {
900 return Some(std::sync::Arc::clone(p));
901 }
902 if *loads
905 >= self
906 .cognitive_config
907 .injection_config
908 .mode_probe_load_budget
909 {
910 return None;
911 }
912 let token = {
913 let mut cache = self.mode_probes.write();
914 match cache.get(&id) {
915 Some(ProbeSlot::Ready(p)) => return Some(std::sync::Arc::clone(p)),
916 Some(ProbeSlot::Loading(_)) => None,
917 None => {
918 let token = NEXT_PROBE_TOKEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
919 cache.insert(id, ProbeSlot::Loading(token));
920 Some(token)
921 }
922 }
923 };
924 *loads += 1;
925 let loaded = self.page_of(id).and_then(|page| {
927 let stamp = self
928 .storage
929 .memory_stamp(page)
930 .filter(|(sid, _)| *sid == id)
931 .map(|(_, s)| (page, s));
932 self.memory_loads
933 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
934 self.storage
935 .load_memory(page)
936 .ok()
937 .filter(|n| n.id == id)
938 .map(|n| std::sync::Arc::new(ModeProbe::from_node(&n, stamp)))
939 });
940 if let Some(token) = token {
941 let mut cache = self.mode_probes.write();
942 if matches!(cache.get(&id), Some(ProbeSlot::Loading(t)) if *t == token) {
943 match &loaded {
944 Some(p) => {
945 cache.insert(id, ProbeSlot::Ready(std::sync::Arc::clone(p)));
946 }
947 None => {
948 cache.remove(&id);
949 }
950 }
951 }
952 }
953 loaded
954 }
955}
956
957pub(crate) enum ProbeSlot {
959 Loading(u64),
960 Ready(std::sync::Arc<ModeProbe>),
961}
962
963static NEXT_PROBE_TOKEN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
964
965pub(crate) const MODE_PROBES_FILE: &str = "mode_probes.bin";
967const PROBES_MAGIC: &[u8; 8] = b"MDBPROB1";
968
969fn put_u64(buf: &mut Vec<u8>, v: u64) {
970 buf.extend_from_slice(&v.to_le_bytes());
971}
972
973fn put_opt(buf: &mut Vec<u8>, v: Option<u64>) {
974 buf.push(u8::from(v.is_some()));
975 put_u64(buf, v.unwrap_or(0));
976}
977
978struct ProbeReader<'a> {
979 buf: &'a [u8],
980 pos: usize,
981}
982
983impl ProbeReader<'_> {
984 fn take(&mut self, n: usize) -> Option<&[u8]> {
985 let end = self.pos.checked_add(n).filter(|&e| e <= self.buf.len())?;
986 let s = &self.buf[self.pos..end];
987 self.pos = end;
988 Some(s)
989 }
990 fn u64(&mut self) -> Option<u64> {
991 Some(u64::from_le_bytes(self.take(8)?.try_into().ok()?))
992 }
993 fn uuid(&mut self) -> Option<uuid::Uuid> {
994 uuid::Uuid::from_slice(self.take(16)?).ok()
995 }
996 fn opt(&mut self) -> Option<Option<u64>> {
997 let flag = self.take(1)?[0];
998 let v = self.u64()?;
999 Some((flag == 1).then_some(v))
1000 }
1001}
1002
1003fn encode_probes(probes: &[(MemoryId, std::sync::Arc<ModeProbe>)]) -> Vec<u8> {
1005 let mut buf = Vec::new();
1006 buf.extend_from_slice(PROBES_MAGIC);
1007 let stamped: Vec<_> = probes
1008 .iter()
1009 .filter_map(|(id, p)| p.stamp.map(|s| (id, p, s)))
1010 .collect();
1011 put_u64(&mut buf, stamped.len() as u64);
1012 for (id, p, (page, stamp)) in stamped {
1013 buf.extend_from_slice(id.0.as_bytes());
1014 put_u64(&mut buf, page.0);
1015 put_u64(&mut buf, stamp);
1016 buf.extend_from_slice(p.agent_id.0.as_bytes());
1017 buf.extend_from_slice(p.user_id.0.as_bytes());
1018 put_opt(&mut buf, p.valid_from);
1019 put_opt(&mut buf, p.valid_until);
1020 put_u64(&mut buf, p.embedding.len() as u64);
1021 for x in &p.embedding {
1022 buf.extend_from_slice(&x.to_le_bytes());
1023 }
1024 }
1025 let crc = crc32fast::hash(&buf);
1026 buf.extend_from_slice(&crc.to_le_bytes());
1027 buf
1028}
1029
1030fn decode_probes(bytes: &[u8]) -> Option<Vec<(MemoryId, ModeProbe)>> {
1032 let (body, crc) = bytes.split_at(bytes.len().checked_sub(4)?);
1033 if crc32fast::hash(body) != u32::from_le_bytes(crc.try_into().ok()?) {
1034 return None;
1035 }
1036 let mut r = ProbeReader { buf: body, pos: 0 };
1037 if r.take(8)? != PROBES_MAGIC {
1038 return None;
1039 }
1040 let n = r.u64()?;
1041 let mut out = Vec::new();
1042 for _ in 0..n {
1043 let id = MemoryId(r.uuid()?);
1044 let page = PageId(r.u64()?);
1045 let stamp = r.u64()?;
1046 let agent_id = AgentId(r.uuid()?);
1047 let user_id = UserId(r.uuid()?);
1048 let valid_from = r.opt()?;
1049 let valid_until = r.opt()?;
1050 let dim = usize::try_from(r.u64()?).ok()?;
1051 let raw = r.take(dim.checked_mul(4)?)?;
1052 let embedding = raw
1053 .as_chunks::<4>()
1054 .0
1055 .iter()
1056 .map(|c| f32::from_le_bytes(*c))
1057 .collect();
1058 out.push((
1059 id,
1060 ModeProbe {
1061 embedding,
1062 agent_id,
1063 user_id,
1064 valid_from,
1065 valid_until,
1066 stamp: Some((page, stamp)),
1067 },
1068 ));
1069 }
1070 (r.pos == body.len()).then_some(out)
1071}
1072
1073impl MenteDb {
1074 pub(crate) fn save_mode_probes(&self) -> MenteResult<()> {
1078 let ready: Vec<(MemoryId, std::sync::Arc<ModeProbe>)> = self
1079 .mode_probes
1080 .read()
1081 .iter()
1082 .filter_map(|(id, slot)| match slot {
1083 ProbeSlot::Ready(p) => Some((*id, std::sync::Arc::clone(p))),
1084 ProbeSlot::Loading(_) => None,
1085 })
1086 .collect();
1087 let bytes = encode_probes(&ready);
1088 let path = self.path.join(MODE_PROBES_FILE);
1089 let tmp = self.path.join(format!("{MODE_PROBES_FILE}.tmp"));
1090 {
1091 use std::io::Write;
1092 let mut f = std::fs::File::create(&tmp)?;
1093 f.write_all(&bytes)?;
1094 f.sync_all()?;
1095 }
1096 std::fs::rename(&tmp, &path)?;
1097 Ok(())
1098 }
1099
1100 pub(crate) fn load_mode_probes(
1105 dir: &std::path::Path,
1106 mapped: &[mentedb_storage::MappedMemory],
1107 ) -> HashMap<MemoryId, ProbeSlot> {
1108 let mut out = HashMap::new();
1109 let Ok(bytes) = std::fs::read(dir.join(MODE_PROBES_FILE)) else {
1110 return out;
1111 };
1112 let Some(probes) = decode_probes(&bytes) else {
1113 tracing::warn!("mode probe snapshot unreadable, probes will load on use");
1114 return out;
1115 };
1116 let current: HashMap<MemoryId, (PageId, u64)> =
1117 mapped.iter().map(|m| (m.id, (m.page, m.stamp))).collect();
1118 for (id, probe) in probes {
1119 if probe.stamp.is_some() && current.get(&id).copied() == probe.stamp {
1120 out.insert(id, ProbeSlot::Ready(std::sync::Arc::new(probe)));
1121 }
1122 }
1123 out
1124 }
1125}
1126
1127#[cfg(test)]
1128mod tests {
1129 use super::*;
1130 use mentedb_core::types::AgentId;
1131 use uuid::Uuid;
1132
1133 #[test]
1134 fn knee_cuts_at_largest_gap() {
1135 let ratio = InjectionConfig::default().knee_gap_ratio;
1136 assert_eq!(knee_cutoff(&[1.0, 0.9, 0.8, 0.1, 0.09], ratio), 3);
1137 assert_eq!(knee_cutoff(&[1.0, 0.9, 0.8, 0.7], ratio), 4);
1139 assert_eq!(knee_cutoff(&[1.0], ratio), 1);
1140 assert_eq!(knee_cutoff(&[], ratio), 0);
1141 }
1142
1143 #[test]
1144 fn attr_counters_roundtrip() {
1145 let mut node = MemoryNode::new(
1146 AgentId(Uuid::nil()),
1147 MemoryType::Semantic,
1148 "x".into(),
1149 vec![0.0; 4],
1150 );
1151 assert_eq!(attr_count(&node, ATTR_INJECTION_SHOWN), 0);
1152 bump_attr(&mut node, ATTR_INJECTION_SHOWN);
1153 bump_attr(&mut node, ATTR_INJECTION_SHOWN);
1154 assert_eq!(attr_count(&node, ATTR_INJECTION_SHOWN), 2);
1155 }
1156
1157 fn node_with_tags(memory_type: MemoryType, tags: &[&str]) -> MemoryNode {
1158 let mut node = MemoryNode::new(AgentId(Uuid::nil()), memory_type, "x".into(), vec![0.0; 4]);
1159 node.tags = tags.iter().map(|t| t.to_string()).collect();
1160 node
1161 }
1162
1163 #[test]
1164 fn tail_gate_trusts_head_and_prunes_weak_tail() {
1165 let cfg = InjectionConfig::default();
1166 let best = 0.8;
1167 assert!(passes_tail_gate(0, 0.05, best, &cfg));
1169 assert!(passes_tail_gate(2, 0.05, best, &cfg));
1170 assert!(passes_tail_gate(3, 0.6, best, &cfg));
1172 assert!(!passes_tail_gate(3, 0.3, best, &cfg));
1173 assert!(passes_tail_gate(5, 0.0, 0.0, &cfg));
1175 }
1176
1177 #[test]
1178 fn cross_project_policy() {
1179 let cfg = InjectionConfig::default();
1180 let best = 0.8;
1181 let foreign_epi = node_with_tags(MemoryType::Episodic, &["scope:project:other"]);
1182 let foreign_sem = node_with_tags(MemoryType::Semantic, &["scope:project:other"]);
1183 let same_epi = node_with_tags(MemoryType::Episodic, &["scope:project:mine"]);
1184 let global_sem = node_with_tags(MemoryType::Semantic, &[]);
1185
1186 assert!(cross_project_allowed(&foreign_epi, 0.1, best, None, &cfg));
1188 assert!(cross_project_allowed(
1190 &same_epi,
1191 0.1,
1192 best,
1193 Some("mine"),
1194 &cfg
1195 ));
1196 assert!(cross_project_allowed(
1197 &global_sem,
1198 0.1,
1199 best,
1200 Some("mine"),
1201 &cfg
1202 ));
1203 assert!(!cross_project_allowed(
1205 &foreign_epi,
1206 0.79,
1207 best,
1208 Some("mine"),
1209 &cfg
1210 ));
1211 assert!(cross_project_allowed(
1213 &foreign_sem,
1214 0.7,
1215 best,
1216 Some("mine"),
1217 &cfg
1218 ));
1219 assert!(!cross_project_allowed(
1220 &foreign_sem,
1221 0.5,
1222 best,
1223 Some("mine"),
1224 &cfg
1225 ));
1226 }
1227
1228 #[test]
1229 fn project_of_reads_scope_tag() {
1230 let node = node_with_tags(MemoryType::Semantic, &["session:s1", "scope:project:apex"]);
1231 assert_eq!(project_of(&node), Some("apex"));
1232 let untagged = node_with_tags(MemoryType::Semantic, &["session:s1"]);
1233 assert_eq!(project_of(&untagged), None);
1234 }
1235}
1236
1237#[cfg(test)]
1238mod mode_probe_tests {
1239 use super::*;
1240 use mentedb_embedding::hash_provider::HashEmbeddingProvider;
1241
1242 #[test]
1243 fn probes_load_once_and_reload_after_a_write() {
1244 let dir = tempfile::tempdir().unwrap();
1245 let db = MenteDb::open_with_embedder(dir.path(), Box::new(HashEmbeddingProvider::new(16)))
1246 .unwrap();
1247 let emb = db.embed_text("rule").unwrap().unwrap();
1248 let mut node = MemoryNode::new(AgentId::nil(), MemoryType::Procedural, "rule".into(), emb);
1249 node.tags = vec!["trigger:x".into()];
1250 let id = node.id;
1251 db.store(node).unwrap();
1252
1253 let mut loads = 0;
1254 assert!(db.mode_probe(id, &mut loads).is_some());
1255 assert!(db.mode_probe(id, &mut loads).is_some());
1256 assert_eq!(loads, 1, "second lookup must come from the cache");
1257
1258 for _ in 0..10 {
1260 db.note_write(MemoryId::new());
1261 }
1262 assert!(db.mode_probe(id, &mut loads).is_some());
1263 assert_eq!(loads, 1, "writes to other memories must not evict it");
1264
1265 db.note_write(id);
1266 assert!(db.mode_probe(id, &mut loads).is_some());
1267 assert_eq!(loads, 2, "a write must force a reload");
1268 assert!(db.mode_probe(id, &mut loads).is_some());
1269 assert_eq!(loads, 2, "and the reload is cached again");
1270 }
1271
1272 #[test]
1276 fn probes_cache_while_other_memories_are_written() {
1277 let dir = tempfile::tempdir().unwrap();
1278 let db = std::sync::Arc::new(
1279 MenteDb::open_with_embedder(dir.path(), Box::new(HashEmbeddingProvider::new(16)))
1280 .unwrap(),
1281 );
1282 let mut ids = Vec::new();
1283 for i in 0..200 {
1284 let text = format!("rule {i}");
1285 let emb = db.embed_text(&text).unwrap().unwrap();
1286 let mut node = MemoryNode::new(AgentId::nil(), MemoryType::Procedural, text, emb);
1287 node.tags = vec!["trigger:x".into()];
1288 ids.push(node.id);
1289 db.store(node).unwrap();
1290 }
1291 let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
1292 let writer = {
1293 let db = std::sync::Arc::clone(&db);
1294 let stop = std::sync::Arc::clone(&stop);
1295 std::thread::spawn(move || {
1296 while !stop.load(std::sync::atomic::Ordering::Relaxed) {
1297 db.note_write(MemoryId::new());
1298 }
1299 })
1300 };
1301 let mut first = 0;
1302 for id in &ids {
1303 db.mode_probe(*id, &mut first);
1304 }
1305 stop.store(true, std::sync::atomic::Ordering::Relaxed);
1306 writer.join().unwrap();
1307 assert_eq!(first, ids.len());
1308 let mut second = 0;
1309 for id in &ids {
1310 db.mode_probe(*id, &mut second);
1311 }
1312 assert_eq!(second, 0, "every probe should be cached after one pass");
1313 }
1314}