From d2fe685ae22d4dfcc6f8126e49ce3da9a2d6fdc5 Mon Sep 17 00:00:00 2001
From: shahondin1624 <throwaway278598@protonmail.com>
Date: Tue, 10 Mar 2026 00:45:46 +0100
Subject: [PATCH] feat: implement provenance tagging and poisoning protection
 (issue #33)

---
 services/memory/src/retrieval/stage3.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/services/memory/src/retrieval/stage3.rs b/services/memory/src/retrieval/stage3.rs
index 3153f12..43491b7 100644
--- a/services/memory/src/retrieval/stage3.rs
+++ b/services/memory/src/retrieval/stage3.rs
@@ -8,6 +8,7 @@ use std::collections::HashMap;
 
 use crate::db::DbError;
 use crate::embedding::store::format_vector_literal;
+use crate::provenance::store as provenance_store;
 use crate::retrieval::RetrievalCandidate;
 use duckdb::Connection;
 
@@ -67,6 +68,14 @@ pub fn load_and_rerank(
     // Compute corpus embedding scores
     let corpus_scores = compute_corpus_scores(conn, query_vector, &id_list)?;
 
+    // Load provenance metadata for all candidates
+    let all_ids: Vec<String> = memories.iter().map(|m| m.id.clone()).collect();
+    let provenance_records = provenance_store::get_provenance_batch(conn, &all_ids)?;
+    let provenance_map: HashMap<String, crate::provenance::ProvenanceRecord> = provenance_records
+        .into_iter()
+        .map(|r| (r.memory_id.clone(), r))
+        .collect();
+
     // Build candidates with all scores
     let mut candidates: Vec<RetrievalCandidate> = memories
         .into_iter()
@@ -91,6 +100,7 @@ pub fn load_and_rerank(
                 description_score: ds,
                 corpus_score: cs,
                 final_score,
+                provenance_metadata: provenance_map.get(&mem.id).cloned(),
             }
         })
         .collect();