Skip to main content

mentedb/
injection.rs

1//! Engine-native injection attention: retrieval shaped for context
2//! injection rather than raw search.
3//!
4//! Plain top-k recall is the wrong contract for injecting memories into a
5//! model's context: the most similar memories to a conversation are echoes
6//! of that conversation, top-k always returns k items however weak the
7//! tail, and near-duplicates (a distilled fact and the turn it came from)
8//! co-rank. This module owns the selection policy that client hooks
9//! previously approximated with heuristics:
10//!
11//! - Session provenance: memories originating from the querying session are
12//!   never returned, they are in that context by construction.
13//! - A working-memory ledger (`exclude_ids`) supplied by the client, holding
14//!   what was already delivered within the current context lifetime.
15//! - A relevance knee instead of a fixed floor: the candidate list is cut at
16//!   the largest score gap, so weak tails vanish without a magic constant.
17//! - Maximal Marginal Relevance over embeddings, so each selected item adds
18//!   information the previous ones did not.
19//! - Type quotas: distilled knowledge first, verbatim episodic capped,
20//!   action notes never (they exist for distillation and session resume).
21//! - User-pinned `scope:always` memories bypass every quality gate.
22//! - Outcome learning: `record_injection_outcome` tracks shown vs used per
23//!   memory; chronically ignored memories are demoted at selection time and
24//!   used ones are reinforced.
25
26use 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
35/// Attribute key: how many times this memory was injected into a context.
36pub const ATTR_INJECTION_SHOWN: &str = "injection_shown";
37/// Attribute key: how many times an injected memory was drawn on by the
38/// reply that followed.
39pub const ATTR_INJECTION_USED: &str = "injection_used";
40
41/// Tunables for injection attention selection.
42#[derive(Debug, Clone)]
43pub struct InjectionConfig {
44    /// Retrieval pool fanned out before selection.
45    pub candidate_pool: usize,
46    /// Cluster completion: when one section tag dominates the selected
47    /// head (at least this fraction of it, minimum two items), the query
48    /// lives inside that cluster, and the lowest ranked non cluster picks
49    /// are swapped for the section's best remaining siblings. A dense
50    /// cluster of sibling rules then arrives together instead of the top
51    /// sibling crowding the rest out. 0 disables. Off by default until the
52    /// offline sweep settles interactions with MMR and the knee and tail
53    /// gates; the mechanism is pinned by tests at 0.34.
54    pub cluster_dominant_fraction: f32,
55    /// Cap on siblings swapped in by cluster completion.
56    pub cluster_fill_max: usize,
57    /// Sections larger than this never count as clusters: a catch all
58    /// section holding half the file dominates any head trivially and its
59    /// siblings are arbitrary, not coherent.
60    pub cluster_max_span: usize,
61    /// Whether injection's hybrid recall includes the lexical (BM25) leg.
62    /// On rule stores the lexical leg can drown the vector signal: generic
63    /// task words match everywhere and rank fusion buries the semantically
64    /// right rule. Off sends vector only retrieval into the same pipeline.
65    pub use_lexical_leg: bool,
66    /// Score multiplier for memories owned by the querying agent. On an
67    /// agent scoped query the agent's own knowledge (its ingested rules,
68    /// its facts) must not be crowded out of the candidate ranking by the
69    /// account's shared nil owned memories; measured on a busy account,
70    /// shared memories otherwise flood the relevant slots. 1.0 disables.
71    pub own_agent_boost: f32,
72    /// Score multiplier for a memory tagged `escalated:violation`: a standing
73    /// rule the agent was shown and then broke. Escalation exists so a rule
74    /// that was violated once surfaces at the FRONT the next time its context
75    /// recurs, not buried mid-list where it was ignored. Without this the
76    /// escalation only bumps salience, which injection does not rank on, so the
77    /// correction would not stick. 1.0 disables.
78    pub escalation_boost: f32,
79    /// The boost applies only to an escalated rule whose cosine is at least
80    /// this fraction of the query's best cosine, so escalation reorders rules
81    /// that are genuinely on topic and never drags an unrelated escalated rule
82    /// into the context on the strength of the boost alone.
83    pub escalation_min_relevance: f32,
84    /// MMR relevance weight; the remainder weighs redundancy.
85    pub mmr_lambda: f32,
86    /// Consecutive score ratio treated as the relevance knee.
87    pub knee_gap_ratio: f32,
88    /// Shown at least this often with zero uses = demoted at selection.
89    pub demotion_shown_min: i64,
90    /// Score multiplier applied to chronically ignored memories.
91    pub demotion_factor: f32,
92    /// Shown at least this often with zero uses = excluded from injection
93    /// entirely (still reachable through ordinary recall). Demotion halves a
94    /// score; this is the terminal stage for memories the model never draws on.
95    pub demotion_drop_shown: i64,
96    /// Head size exempt from the semantic tail gate. The top RRF ranks are
97    /// trusted as-is because a pure keyword hit (exact identifier, phone
98    /// number) can carry a low cosine yet be exactly right.
99    pub tail_min_keep: usize,
100    /// Beyond the head, a candidate must reach this fraction of the best
101    /// query cosine among candidates or it is dropped. Relative to the query's
102    /// own score distribution, never an absolute cosine, so it holds across
103    /// embedders. This is what keeps a weak tail silent instead of letting
104    /// selection fill its quota with plausible-but-irrelevant memories.
105    pub tail_floor_fraction: f32,
106    /// Whether episodic memories tagged to another project may inject when a
107    /// current project is set. Episodic turns are working context of the
108    /// project they happened in; injecting them elsewhere is the classic
109    /// cross-context contamination users disable memory over.
110    pub cross_project_episodic: bool,
111    /// A non-episodic memory tagged to another project must reach this
112    /// fraction of the best query cosine to inject cross-project. Facts and
113    /// preferences may travel, but only when they are among the most relevant
114    /// candidates for the query, not as tail filler.
115    pub cross_project_floor_fraction: f32,
116    /// Reply similarity above which a shown memory counts as used.
117    pub used_similarity: f32,
118    /// Salience reinforcement applied when a memory is actually used.
119    pub use_reinforcement: f32,
120    /// Associative recall: how many 1-hop graph neighbors of the top vector hits
121    /// to fold into the candidate pool, so a memory linked to a strong hit can
122    /// surface even when it is not itself vector-similar to the query. 0 (the
123    /// default) disables expansion, leaving pure vector/hybrid recall unchanged.
124    pub graph_expansion_max: usize,
125    /// Score a neighbor inherits from the hit it was reached through, as a
126    /// fraction of that hit's score times the edge weight. Keeps neighbors below
127    /// direct hits so they only win a slot when there is room after them.
128    pub graph_expansion_decay: f32,
129    /// Mode rules: standing directives bound to an activity (working on
130    /// code, composing a reply) carry `trigger:` tags, and a trigger is
131    /// active for a turn when the turn's embedding lands within this cosine
132    /// of one of that trigger's own rules. Self defining and cognitive,
133    /// the rules are their own anchors, no pattern matching anywhere.
134    /// Embedder coupled like every absolute gate. 0 disables the channel.
135    pub mode_trigger_gate: f32,
136    /// Maximum mode rules injected per turn across all active triggers.
137    pub mode_rules_max: usize,
138    /// Most mode probes one injection recall may load from storage on cache
139    /// misses. Probes persist across reopen, so a miss is normally a rule
140    /// written since the last snapshot; this caps the cost when the persisted
141    /// probes are missing (first open after an upgrade, a deleted file), where
142    /// a large rule set would otherwise load every rule inside one request.
143    /// Rules past the cap are skipped for that request and load on later ones.
144    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
178/// A request for injection-ready context.
179pub struct InjectionQuery<'a> {
180    /// Embedding of the prompt (plus any conversational blend).
181    pub embedding: &'a [f32],
182    /// Raw query text for the lexical half of hybrid recall.
183    pub query_text: Option<&'a str>,
184    /// The session issuing the query. Memories tagged `session:<id>` with
185    /// this ID are excluded: they are already in that session's context.
186    pub session_id: Option<&'a str>,
187    /// Memory IDs already delivered within the current context lifetime.
188    pub exclude_ids: &'a [MemoryId],
189    /// Maximum items returned beyond pinned memories.
190    pub max_items: usize,
191    /// Maximum verbatim episodic items within `max_items`.
192    pub max_episodic: usize,
193    /// Restrict recall to this agent's memories plus shared (nil owned)
194    /// knowledge. None recalls globally.
195    pub agent_id: Option<AgentId>,
196    /// Restrict recall to this user's memories plus shared (nil owned)
197    /// knowledge, orthogonal to `agent_id`. None recalls globally on the user
198    /// axis. A memory is injectable only when visible on BOTH axes.
199    pub user_id: Option<UserId>,
200    /// Current project (a `scope:project:<name>` value without the prefix). When
201    /// set, memories from other projects are weighted down in recall so injection
202    /// favors the active project instead of dumping cross-project noise. None
203    /// recalls across all projects.
204    pub current_project: Option<&'a str>,
205}
206
207/// Why an item was selected, for introspection and client display.
208#[derive(Debug, Clone, Copy, PartialEq, Eq)]
209pub enum SelectionReason {
210    /// User-pinned scope:always memory: bypasses every quality gate.
211    Pinned,
212    /// Survived the relevance knee, MMR, and quotas.
213    Relevant,
214    /// A standing directive whose activity mode is active for this turn:
215    /// the turn's embedding landed within the mode gate of the rule.
216    ModeRule,
217}
218
219/// One injection-ready memory with its selection metadata.
220pub struct InjectionCandidate {
221    pub node: MemoryNode,
222    /// Retrieval score after demotion adjustment (0.0 for pinned items,
223    /// which are not score-ranked).
224    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
265/// Cut a descending score list at its largest relative gap (the relevance
266/// knee), when that gap is decisive. Returns how many items to keep.
267fn 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
289/// The project a memory is scoped to, from its `scope:project:<name>` tag.
290fn project_of(node: &MemoryNode) -> Option<&str> {
291    node.tags
292        .iter()
293        .find_map(|t| t.strip_prefix("scope:project:"))
294}
295
296/// Cross-project policy for a candidate when a current project is set.
297/// Untagged (global) memories and same-project memories always pass. Episodic
298/// memories from another project pass only if configured to. Other types from
299/// another project must be among the most semantically relevant candidates
300/// (cosine relative to the query's best), never tail filler.
301fn 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
323/// Semantic tail gate over a descending-ranked candidate list: the head
324/// (`tail_min_keep`) is trusted as ranked, the tail must reach
325/// `tail_floor_fraction` of the best query cosine to stay. Operates on
326/// (index, cosine) so RRF rank compression does not blunt it. Degenerate
327/// inputs (best cosine 0, e.g. mid-migration dimension mismatch) pass
328/// everything rather than silencing injection entirely.
329fn 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    /// Injection-ready context selection. See the module docs for the
338    /// policy; this is the API client hooks should prefer over raw recall.
339    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        // Phase timings, logged when a call is slow, so a slow recall in
348        // production names the phase and how many memories it loaded.
349        let t_start = std::time::Instant::now();
350        // Candidate pool from hybrid recall.
351        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                // Actions never inject; pinned items are handled separately.
382                // Ghost memories are speculative working material for the
383                // trajectory tracker, not confirmed knowledge; injecting
384                // "Unconfirmed: ..." as if it were a fact misleads.
385                continue;
386            }
387            if let Some(ref st) = session_tag
388                && has_tag(&node, st)
389            {
390                continue;
391            }
392            // Chronically ignored memories fall back in the ranking until
393            // usage or decay resolves them; past the drop threshold they are
394            // excluded from injection outright (ordinary recall still finds
395            // them). Exposure without a single use is the strongest negative
396            // signal the system has.
397            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        // Pre gate snapshot of section tagged candidates: cluster completion
418        // may need siblings the knee and tail gates are about to drop.
419        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        // Gates are relative to the query's own best cosine, never an absolute
430        // threshold, so they hold across embedders (RRF scores are rank
431        // compressed and useless for this; the raw cosine is not).
432        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        // Semantic tail gate: beyond the trusted head, weak-cosine candidates
442        // vanish instead of being carried to MMR, which would otherwise fill
443        // max_items however irrelevant the tail. Silence beats plausible noise.
444        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        // Escalation re-ranks AMONG the relevance survivors: a rule the agent
452        // was shown and then broke leads the next time its context recurs. It
453        // runs here, after the cosine gates, so the boost can never drag an
454        // irrelevant escalated rule into the context, only reorder rules that
455        // already earned their place.
456        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        // Associative recall: after the vector-tail knee, fold in 1-hop graph
474        // neighbors of the surviving hits, so a memory linked to a relevant hit
475        // can surface even when it is not itself vector-similar to the query.
476        // Runs after the knee (which is for weak vector tails, a different signal
477        // than an explicit edge) so a neighbor's decayed score is not mistaken
478        // for a tail and cut. Disabled by default (graph_expansion_max == 0),
479        // leaving pure vector/hybrid recall intact. Neighbors inherit a decayed
480        // share of the hit's score so they never outrank direct hits, pass the
481        // same gates, and are bounded by graph_expansion_max; MMR then picks the
482        // final set within max_items.
483        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        // Cluster dominance is read off the score ranking BEFORE MMR: MMR
527        // deliberately de clusters the head, hiding the very signal cluster
528        // completion needs.
529        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            // A section wider than cluster_max_span across the whole pool is
540            // a catch all, not a coherent cluster; it dominates any head
541            // trivially and its siblings are arbitrary.
542            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        // MMR selection with type quotas.
564        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        // Cluster completion: when one section dominates the selected head,
607        // the turn lives inside that cluster of sibling rules, and embedding
608        // rank underestimates the rest of the cluster (the harness measured
609        // exactly this on a real 22KB agent file: two governing rules ranked
610        // below k=24 behind their own section siblings). Swap the weakest
611        // non cluster tail picks for the section's best remaining siblings,
612        // gates overridden deliberately, bounded by cluster_fill_max.
613        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                // Evict the weakest non cluster picks, never the top pick.
628                // The span cap above keeps this to coherent clusters, where
629                // the cluster hypothesis says sibling relevance is
630                // underestimated by embedding rank.
631                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        // Pinned memories bypass every gate except the ledger, and lead the
651        // result. The client's ledger reset (at context loss) governs
652        // re-delivery. Fetched from the tag index (scope:always is capped
653        // near-zero) so this loads only the handful of matches; the old path
654        // scanned every page on every injection call, which thrashed the
655        // buffer pool against disk on large accounts.
656        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        // Mode rules: standing directives bound to an activity rather than a
676        // topic or a tool moment. "Longterm fixes only" must fire for "fix
677        // the login timeout" even though they share no words; the turn's
678        // embedding landing near the rule's own stored embedding is the
679        // activation, cognitive and self defining, dot products only.
680        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                // Activation signal: the trigger's exemplar turns when it
691                // has any (model emitted examples of requests that enter
692                // this mode; a turn matching a turn is the strong signal,
693                // measured), otherwise the rules' own embeddings.
694                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                // The mode is active: its RULES (never exemplars) become
726                // candidates. Validity and owner come from the cached probes;
727                // full memories are loaded below only for the few kept.
728                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            // Strongest mode first (stable, so rules keep their index order
749            // within a mode), then load only until the cap is filled. A rule
750            // whose stored tags no longer carry the trigger is skipped, as
751            // before.
752            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    /// Record the outcome of an injection: every shown memory's exposure
794    /// count rises, and memories the reply actually drew on (embedding
795    /// similarity against the reply) are counted as used and reinforced.
796    /// Returns (shown_updated, used_count).
797    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            // Retrieval reinforcement: being surfaced at all is an access. Refresh
821            // the decay clock and bump the access count for every shown memory, not
822            // only the ones the reply echoed, so anything that keeps getting
823            // recalled stays alive instead of decaying to the forget threshold while
824            // it is actively in use. Memories that are never recalled still get no
825            // touch here, so they decay and are forgotten as intended; only what
826            // retrieval keeps surfacing survives.
827            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                // The reply actually drew on it: a stronger signal than mere
835                // exposure, so add a salience boost on top of the access refresh.
836                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            // Counter updates must not re-run write inference; write the
842            // node in place.
843            self.update_node(pid, &node)?;
844            updated += 1;
845        }
846        Ok((updated, used_total))
847    }
848}
849
850/// What mode activation needs from a `trigger:` rule or `mode-exemplar:` turn:
851/// its embedding, validity window and owner. Cached per memory on the
852/// database so activation never loads the full memory from storage.
853pub(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    /// The memory's page and storage write stamp, read BEFORE the memory was
860    /// loaded. The probe may be newer than the stamp, never older, so a
861    /// reopen that finds the same page and stamp knows no write landed on the
862    /// memory since and can trust the probe without loading it. None when
863    /// storage had no stamp (the probe is then never persisted).
864    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    /// Same rule as `MemoryNode::is_valid_at`.
880    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    /// The cached mode probe for `id`, loading the memory once on a miss
891    /// (counted in `loads`). The slot is marked pending while loading; a write
892    /// to this memory removes the slot, so a copy read before that write is
893    /// returned but never cached. Writes to other memories do not interfere.
894    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        // Past the per request budget a miss is skipped, not loaded: it loads
903        // on a later request, so one request never pays for every rule.
904        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        // Stamp first, then load: see `ModeProbe::stamp`.
926        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
957/// A mode probe cache slot: being loaded (by the holder of this token) or ready.
958pub(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
965/// File (in the database directory) holding the persisted mode probes.
966pub(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
1003/// Encode stamped probes: magic, count, records, CRC32 of everything before.
1004fn 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
1030/// Decode a probe file; None when it is malformed in any way.
1031fn 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    /// Persist the ready mode probes, each with the page and write stamp it
1075    /// was loaded at, so a reopen restores them without loading the rules.
1076    /// Atomic (temp file, fsync, rename).
1077    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    /// The persisted probes still current for `mapped` (the memory map just
1101    /// loaded on open): a probe is restored only when its memory is on the
1102    /// same page with the same write stamp, i.e. nothing wrote the memory
1103    /// since the probe was loaded. Anything else is dropped and loads on use.
1104    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        // No decisive gap: keep everything.
1138        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        // Head positions pass regardless of cosine (keyword hits).
1168        assert!(passes_tail_gate(0, 0.05, best, &cfg));
1169        assert!(passes_tail_gate(2, 0.05, best, &cfg));
1170        // Tail passes only near the query's own best cosine.
1171        assert!(passes_tail_gate(3, 0.6, best, &cfg));
1172        assert!(!passes_tail_gate(3, 0.3, best, &cfg));
1173        // Degenerate best cosine (mid-migration): everything passes.
1174        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        // No current project: everything passes.
1187        assert!(cross_project_allowed(&foreign_epi, 0.1, best, None, &cfg));
1188        // Same project and untagged always pass.
1189        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        // Foreign episodic never passes by default.
1204        assert!(!cross_project_allowed(
1205            &foreign_epi,
1206            0.79,
1207            best,
1208            Some("mine"),
1209            &cfg
1210        ));
1211        // Foreign semantic passes only when among the most relevant.
1212        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        // Writes to other memories leave this probe cached.
1259        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    /// The production failure: an account writing constantly (a sync backlog
1273    /// draining) while recall fills the cache. Writes to other memories during
1274    /// a load must not stop that probe from being cached.
1275    #[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}