Fixes to locking mechanism

src/index.rs

@@ -870,6 +870,51 @@ impl SemanticIndex {
         }
     }
 
+    /// M29 two-level conversation lookup under a **single** read lock.
+    ///
+    /// Replaces the double-lock pattern in `process_query_locally` where the
+    /// handler acquired `state.index.read()` once for the conversation namespace
+    /// and then — if that missed — released the lock and re-acquired it for
+    /// the base namespace fallback. That open window between lock 1 and lock 2
+    /// allowed a concurrent write to alter index state, meaning the two reads
+    /// observed inconsistent snapshots of the index (TOCTOU).
+    ///
+    /// By running both namespace lookups inside the same `&self` borrow (the
+    /// caller holds the `RwLock` read guard for the duration), this function
+    /// gives a single consistent snapshot: if the conv namespace is empty now,
+    /// the base namespace result is from the same instant, not a later one.
+    ///
+    /// Returns `(hit, hit_namespace, scope)`:
+    /// - `hit` — the `QueryHit` if found, `None` on miss
+    /// - `hit_namespace` — the namespace that produced the result; always set
+    ///   so the caller can pass it directly to `record_access` and metrics
+    /// - `scope` — `Some("conversation")` | `Some("global")` | `None` (miss)
+    pub fn query_two_level(
+        &self,
+        embedding: &[f32],
+        threshold: f32,
+        conv_ns: &str,
+        base_ns: &str,
+        query_text: Option<&str>,
+        now: u64,
+    ) -> Result<(Option<QueryHit>, String, Option<&'static str>)> {
+        // Conversation namespace has priority.
+        if let Some(ns) = self.namespaces.get(conv_ns) {
+            if let Some(hit) = ns.query_with_text(embedding, threshold, query_text, now)? {
+                return Ok((Some(hit), conv_ns.to_string(), Some("conversation")));
+            }
+        }
+        // Fallback: base namespace.
+        match self.namespaces.get(base_ns) {
+            Some(ns) => {
+                let hit = ns.query_with_text(embedding, threshold, query_text, now)?;
+                let scope = if hit.is_some() { Some("global") } else { None };
+                Ok((hit, base_ns.to_string(), scope))
+            }
+            None => Ok((None, base_ns.to_string(), None)),
+        }
+    }
+
     pub fn entry_count(&self) -> usize {
         self.namespaces.values().map(|n| n.entry_count()).sum()
     }
@@ -2242,4 +2287,138 @@ mod tests {
         );
         assert!(idx.namespaces.contains_key(&conv_ns));
     }
+
+    // --- query_two_level unit tests ----------------------------------------
+
+    #[test]
+    fn test_query_two_level_conv_hit() {
+        // Entry in conv namespace → returns scope="conversation", conv ns name.
+        let mut idx = test_index();
+        let base = base_namespace("m::3", None);
+        let conv = conversation_namespace("m::3", None, "s1");
+        idx.insert(vec![1.0_f32, 0.0, 0.0], "base-r".into(), "q".into(), &base)
+            .unwrap();
+        idx.insert(vec![1.0_f32, 0.0, 0.0], "conv-r".into(), "q".into(), &conv)
+            .unwrap();
+        let now = now_unix_secs();
+        let (hit, ns, scope) = idx
+            .query_two_level(&[1.0_f32, 0.0, 0.0], 0.9, &conv, &base, None, now)
+            .unwrap();
+        assert_eq!(hit.unwrap().response, "conv-r");
+        assert_eq!(ns, conv);
+        assert_eq!(scope, Some("conversation"));
+    }
+
+    #[test]
+    fn test_query_two_level_base_fallback() {
+        // Conv namespace is empty → falls through to base, returns scope="global".
+        let mut idx = test_index();
+        let base = base_namespace("m::3", None);
+        let conv = conversation_namespace("m::3", None, "s1");
+        idx.insert(vec![1.0_f32, 0.0, 0.0], "base-r".into(), "q".into(), &base)
+            .unwrap();
+        let now = now_unix_secs();
+        let (hit, ns, scope) = idx
+            .query_two_level(&[1.0_f32, 0.0, 0.0], 0.9, &conv, &base, None, now)
+            .unwrap();
+        assert_eq!(hit.unwrap().response, "base-r");
+        assert_eq!(ns, base);
+        assert_eq!(scope, Some("global"));
+    }
+
+    #[test]
+    fn test_query_two_level_both_miss() {
+        // Nothing in either namespace → None, hit_ns is base, scope is None.
+        let idx = test_index();
+        let base = base_namespace("m::3", None);
+        let conv = conversation_namespace("m::3", None, "s1");
+        let now = now_unix_secs();
+        let (hit, ns, scope) = idx
+            .query_two_level(&[1.0_f32, 0.0, 0.0], 0.9, &conv, &base, None, now)
+            .unwrap();
+        assert!(hit.is_none());
+        assert_eq!(ns, base);
+        assert_eq!(scope, None);
+    }
+
+    #[test]
+    fn test_query_two_level_conv_priority_over_base() {
+        // Distinct responses in both namespaces: conv always wins when both
+        // have the vector, regardless of insertion order.
+        let mut idx = test_index();
+        let base = base_namespace("m::3", None);
+        let conv = conversation_namespace("m::3", None, "s1");
+        idx.insert(
+            vec![1.0_f32, 0.0, 0.0],
+            "from-base".into(),
+            "q".into(),
+            &base,
+        )
+        .unwrap();
+        idx.insert(
+            vec![0.8_f32, 0.6, 0.0],
+            "from-conv".into(),
+            "q2".into(),
+            &conv,
+        )
+        .unwrap();
+        let now = now_unix_secs();
+        let (hit, _, scope) = idx
+            .query_two_level(&[1.0_f32, 0.0, 0.0], 0.5, &conv, &base, None, now)
+            .unwrap();
+        // conv has a nearby vector; base has an exact match — conv is checked
+        // first so "from-conv" is returned.
+        assert_eq!(hit.unwrap().response, "from-conv");
+        assert_eq!(scope, Some("conversation"));
+    }
+
+    #[test]
+    fn test_query_two_level_exact_match_in_conv() {
+        // Exact-match pre-filter works through query_two_level: text match in
+        // conv namespace hits even with an orthogonal embedding.
+        let mut idx = test_index();
+        let base = base_namespace("m::3", None);
+        let conv = conversation_namespace("m::3", None, "s1");
+        idx.insert(
+            vec![1.0_f32, 0.0, 0.0],
+            "Paris".into(),
+            "Capital of France?".into(),
+            &conv,
+        )
+        .unwrap();
+        let now = now_unix_secs();
+        let (hit, _, scope) = idx
+            .query_two_level(
+                &[0.0_f32, 0.0, 1.0], // orthogonal — HNSW would miss
+                0.99,
+                &conv,
+                &base,
+                Some("Capital of France?"),
+                now,
+            )
+            .unwrap();
+        let h = hit.unwrap();
+        assert!(h.exact_match);
+        assert_eq!(h.response, "Paris");
+        assert_eq!(scope, Some("conversation"));
+    }
+
+    #[test]
+    fn test_query_two_level_error_propagates() {
+        // Dimension mismatch in the conv namespace propagates as Err rather
+        // than silently falling through to the base namespace.
+        let mut idx = test_index();
+        let base = base_namespace("m::3", None);
+        let conv = conversation_namespace("m::3", None, "s1");
+        // Seed conv with a 3-d entry so the namespace knows its dimension.
+        idx.insert(vec![1.0_f32, 0.0, 0.0], "r".into(), "q".into(), &conv)
+            .unwrap();
+        let now = now_unix_secs();
+        // Query with a 4-d vector — dimension mismatch inside conv namespace.
+        let result = idx.query_two_level(&[1.0_f32, 0.0, 0.0, 0.0], 0.9, &conv, &base, None, now);
+        assert!(
+            result.is_err(),
+            "dimension mismatch must propagate as Err, not fall through to base"
+        );
+    }
 }

src/server.rs

@@ -231,54 +231,43 @@ async fn process_query_locally(
     let conv_ns =
         conv_id.map(|cid| conversation_namespace(model_id, req.cache_scope.as_deref(), cid));
 
-    // First-level lookup: conversation namespace if present, else base.
-    let primary_ns: &str = conv_ns.as_deref().unwrap_or(&base);
-    let primary_result = {
+    // Acquire the index read lock exactly once — a single consistent snapshot
+    // covers both the conversation namespace check and the base fallback.
+    // The previous double-lock pattern released and re-acquired between the
+    // two lookups, opening a TOCTOU window where a concurrent WAL flush could
+    // insert into the conversation namespace and be missed, or where an
+    // expiry/eviction between the two acquisitions could yield an inconsistent
+    // result. query_two_level runs both lookups under the same &self borrow,
+    // eliminating the window and halving the lock-acquire count on the hot path.
+    let (final_result, hit_ns_owned, hit_scope): (
+        Result<Option<crate::index::QueryHit>, anyhow::Error>,
+        String,
+        Option<&'static str>,
+    ) = if let Some(ref cns) = conv_ns {
         let index = state.index.read().await;
-        index.query_with_text(
+        match index.query_two_level(
             &req.embedding,
             req.threshold,
-            primary_ns,
+            cns,
+            &base,
             req.query_text.as_deref(),
             now,
-        )
-    };
-
-    // Fallback (only when we had a conversation_id and the conversation
-    // namespace missed). Errors from the primary still propagate as 400.
-    let (final_result, hit_ns_owned, hit_scope): (
-        Result<Option<crate::index::QueryHit>, anyhow::Error>,
-        String,
-        Option<&'static str>,
-    ) = match primary_result {
-        Ok(Some(hit)) => {
-            let scope = if conv_ns.is_some() {
-                Some("conversation")
-            } else {
-                None
-            };
-            (Ok(Some(hit)), primary_ns.to_string(), scope)
+        ) {
+            Ok((hit, ns, scope)) => (Ok(hit), ns, scope),
+            Err(e) => (Err(e), base.clone(), None),
         }
-        Ok(None) if conv_ns.is_some() => {
-            // Two-level fallback: try the base namespace.
-            let fallback = {
-                let index = state.index.read().await;
-                index.query_with_text(
-                    &req.embedding,
-                    req.threshold,
-                    &base,
-                    req.query_text.as_deref(),
-                    now,
-                )
-            };
-            match fallback {
-                Ok(Some(hit)) => (Ok(Some(hit)), base.clone(), Some("global")),
-                Ok(None) => (Ok(None), base.clone(), None),
-                Err(e) => (Err(e), base.clone(), None),
-            }
+    } else {
+        let index = state.index.read().await;
+        match index.query_with_text(
+            &req.embedding,
+            req.threshold,
+            &base,
+            req.query_text.as_deref(),
+            now,
+        ) {
+            Ok(hit) => (Ok(hit), base.clone(), None),
+            Err(e) => (Err(e), base.clone(), None),
         }
-        Ok(None) => (Ok(None), primary_ns.to_string(), None),
-        Err(e) => (Err(e), primary_ns.to_string(), None),
     };
 
     let elapsed = start.elapsed().as_secs_f64();