diff --git a/src/adapters/neo4j.rs b/src/adapters/neo4j.rs index c8bf875..03b9dd9 100644 --- a/src/adapters/neo4j.rs +++ b/src/adapters/neo4j.rs @@ -89,7 +89,7 @@ impl Neo4jPort for Neo4jAdapter { self.cypher( statement, json!({ - "tenant_id": "configured-by-user", + "tenant_id": node.tenant_id, "user_id": node.user_id, "label": node.label, "canonical_name": node.canonical_name, @@ -116,7 +116,7 @@ impl Neo4jPort for Neo4jAdapter { self.cypher( statement, json!({ - "tenant_id": "configured-by-user", + "tenant_id": edge.tenant_id, "user_id": edge.user_id, "from_key": edge.from_key, "to_key": edge.to_key, diff --git a/src/adapters/redis.rs b/src/adapters/redis.rs index cc251f4..f49cc63 100644 --- a/src/adapters/redis.rs +++ b/src/adapters/redis.rs @@ -63,15 +63,26 @@ impl RedisPort for RedisAdapter { fn invalidate_prefix(&self, prefix: &str) -> CoreResult<()> { let mut connection = self.connection()?; let pattern = self.namespaced_key(&format!("{prefix}*")); - let keys: Vec = redis::cmd("KEYS") - .arg(pattern) - .query(&mut connection) - .map_err(|error| CoreError::Io(error.to_string()))?; - if !keys.is_empty() { - let _: () = redis::cmd("DEL") - .arg(keys) + let mut cursor: u64 = 0; + loop { + let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") + .arg(cursor) + .arg("MATCH") + .arg(&pattern) + .arg("COUNT") + .arg(100) .query(&mut connection) .map_err(|error| CoreError::Io(error.to_string()))?; + if !keys.is_empty() { + let _: () = redis::cmd("DEL") + .arg(keys) + .query(&mut connection) + .map_err(|error| CoreError::Io(error.to_string()))?; + } + cursor = next_cursor; + if cursor == 0 { + break; + } } Ok(()) } @@ -92,16 +103,19 @@ impl RedisPort for RedisAdapter { fn release_lease(&self, key: &str, owner: &str) -> CoreResult<()> { let mut connection = self.connection()?; let namespaced = self.namespaced_key(key); - let current: Option = redis::cmd("GET") - .arg(&namespaced) - .query(&mut connection) + // Atomic check-and-delete using Lua script to prevent TOCTOU race condition + let script = redis::Script::new( + r#"if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) + else + return 0 + end"#, + ); + let _: i32 = script + .key(&namespaced) + .arg(owner) + .invoke(&mut connection) .map_err(|error| CoreError::Io(error.to_string()))?; - if current.as_deref() == Some(owner) { - let _: () = redis::cmd("DEL") - .arg(namespaced) - .query(&mut connection) - .map_err(|error| CoreError::Io(error.to_string()))?; - } Ok(()) } } diff --git a/src/domain/graph.rs b/src/domain/graph.rs index 3017c18..ece48d9 100644 --- a/src/domain/graph.rs +++ b/src/domain/graph.rs @@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct GraphNode { + pub tenant_id: String, pub user_id: String, pub label: String, pub name: String, @@ -18,6 +19,7 @@ pub struct GraphNode { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct GraphEdge { + pub tenant_id: String, pub user_id: String, pub from_key: String, pub relationship_type: String, @@ -62,10 +64,11 @@ pub struct GraphifyOutput { } impl GraphNode { - pub fn new(user_id: &str, label: &str, name: &str, confidence: f32) -> CoreResult { + pub fn new(tenant_id: &str, user_id: &str, label: &str, name: &str, confidence: f32) -> CoreResult { validate_score("confidence", confidence)?; let canonical_name = canonicalize(name); Ok(Self { + tenant_id: tenant_id.to_string(), user_id: user_id.to_string(), label: label.to_string(), name: name.trim().to_string(), @@ -79,6 +82,7 @@ impl GraphNode { impl GraphEdge { pub fn new( + tenant_id: &str, user_id: &str, from_key: &str, relationship_type: &str, @@ -89,6 +93,7 @@ impl GraphEdge { validate_score("confidence", confidence)?; let now = now_timestamp(); Ok(Self { + tenant_id: tenant_id.to_string(), user_id: user_id.to_string(), from_key: from_key.to_string(), relationship_type: relationship_type.trim().to_uppercase(), @@ -106,7 +111,7 @@ pub fn graphify_record(record: &MemoryRecord, hints: &[GraphHint]) -> CoreResult let mut contradictions = Vec::new(); for entity in &record.entities { if !entity.trim().is_empty() { - nodes.push(GraphNode::new(&record.user_id, infer_label(entity), entity, 0.8)?); + nodes.push(GraphNode::new(&record.tenant_id, &record.user_id, infer_label(entity), entity, 0.8)?); } } @@ -116,18 +121,21 @@ pub fn graphify_record(record: &MemoryRecord, hints: &[GraphHint]) -> CoreResult for hint in &all_hints { validate_graph_hint(hint)?; let from = GraphNode::new( + &record.tenant_id, &record.user_id, &hint.from_label, &hint.from_name, hint.confidence, )?; let to = GraphNode::new( + &record.tenant_id, &record.user_id, &hint.to_label, &hint.to_name, hint.confidence, )?; relationships.push(GraphEdge::new( + &record.tenant_id, &record.user_id, &from.key, &hint.relationship_type, @@ -177,10 +185,10 @@ pub fn canonicalize(name: &str) -> String { name.trim().to_lowercase().replace(' ', "_") } -pub fn connect(edges: &mut Vec, from: impl Into, to: impl Into) { +pub fn connect(edges: &mut Vec, tenant_id: &str, from: impl Into, to: impl Into) { let from = from.into(); let to = to.into(); - if let Ok(edge) = GraphEdge::new("default", &from, "RELATED_TO", &to, 0.5, "legacy") { + if let Ok(edge) = GraphEdge::new(tenant_id, "default", &from, "RELATED_TO", &to, 0.5, "legacy") { edges.push(edge); } } diff --git a/src/domain/memory.rs b/src/domain/memory.rs index 7ec6d7b..5a29c52 100644 --- a/src/domain/memory.rs +++ b/src/domain/memory.rs @@ -254,15 +254,26 @@ pub fn validate_score(name: &str, score: f32) -> CoreResult<()> { pub fn now_timestamp() -> String { SystemTime::now() .duration_since(UNIX_EPOCH) - .unwrap_or_default() + .expect("System time is before UNIX epoch (1970-01-01). Check system clock.") .as_secs() .to_string() } pub fn deterministic_id(parts: &[&str]) -> String { - let mut hasher = DefaultHasher::new(); - parts.hash(&mut hasher); - format!("mem_{:016x}", hasher.finish()) + // Use FNV-1a hash which is stable across Rust versions and has good distribution + const FNV_OFFSET: u64 = 0xcbf29ce484222325; + const FNV_PRIME: u64 = 0x100000001b3; + let mut hash = FNV_OFFSET; + for part in parts { + for byte in part.as_bytes() { + hash ^= *byte as u64; + hash = hash.wrapping_mul(FNV_PRIME); + } + // Add separator to prevent "ab"+"c" == "a"+"bc" collisions + hash ^= 0xFF; + hash = hash.wrapping_mul(FNV_PRIME); + } + format!("mem_{:016x}", hash) } pub fn estimate_tokens(text: &str) -> u32 { diff --git a/src/lib.rs b/src/lib.rs index b6a1dca..750c66c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -423,7 +423,7 @@ mod tests { ) .unwrap(); assert!(response.telemetry.vector_candidates >= 1); - assert!(response.telemetry.vector_ms <= response.telemetry.vector_ms + 1); + assert!(response.telemetry.vector_ms < 10000, "vector search took too long: {}ms", response.telemetry.vector_ms); assert!(response.telemetry.token_utilization >= 0.0); assert!(response.telemetry.dedupe_ratio >= 0.0); } diff --git a/src/package/mod.rs b/src/package/mod.rs index 9fcabbb..8595c0f 100644 --- a/src/package/mod.rs +++ b/src/package/mod.rs @@ -117,6 +117,11 @@ fn control_plane() -> &'static Mutex { INSTANCE.get_or_init(|| Mutex::new(RuntimeControlPlane::default())) } +fn shared_store() -> &'static Mutex { + static INSTANCE: OnceLock> = OnceLock::new(); + INSTANCE.get_or_init(|| Mutex::new(TestMemoryStore::new())) +} + pub fn e2e_smoke_json() -> Result { let mut store = TestMemoryStore::new(); let appended = append_session_message( @@ -291,6 +296,7 @@ pub fn adapter_smoke_json(request_json: &str) -> Result { }, )?; neo4j.merge_node(&crate::graph::GraphNode::new( + &request.tenant_id, &request.user_id, "Entity", &request.memory_id, @@ -369,14 +375,18 @@ pub fn mcp_call_json(request_json: &str) -> Result { "nextral.memory.forget" => { let payload: ForgetMemoryRequest = serde_json::from_str(&request.payload_json).map_err(CoreError::from)?; - let mut store = TestMemoryStore::new(); - Ok(serde_json::to_string(&forget_memory(&mut store, payload)?).map_err(CoreError::from)?) + let mut store = shared_store() + .lock() + .map_err(|error| CoreError::Conflict(error.to_string()))?; + Ok(serde_json::to_string(&forget_memory(&mut *store, payload)?).map_err(CoreError::from)?) } "nextral.reminders.due" => { let payload: ExecuteDueRemindersRequest = serde_json::from_str(&request.payload_json).map_err(CoreError::from)?; - let mut store = TestMemoryStore::new(); - Ok(serde_json::to_string(&execute_due_reminders(&mut store, payload)?).map_err(CoreError::from)?) + let mut store = shared_store() + .lock() + .map_err(|error| CoreError::Conflict(error.to_string()))?; + Ok(serde_json::to_string(&execute_due_reminders(&mut *store, payload)?).map_err(CoreError::from)?) } "experiments.create" => { let payload: ExperimentCreateRequest = @@ -481,8 +491,10 @@ pub fn mcp_call_json(request_json: &str) -> Result { "nextral.graph.query" => { let payload: RetrievalRequest = serde_json::from_str(&request.payload_json).map_err(CoreError::from)?; - let mut store = TestMemoryStore::new(); - let response = crate::runtime::retrieval::retrieve(&mut store, payload)?; + let mut store = shared_store() + .lock() + .map_err(|error| CoreError::Conflict(error.to_string()))?; + let response = crate::runtime::retrieval::retrieve(&mut *store, payload)?; let graph_only: Vec<_> = response .items .into_iter() diff --git a/src/runtime/prospective.rs b/src/runtime/prospective.rs index c2022ef..9bd7dc2 100644 --- a/src/runtime/prospective.rs +++ b/src/runtime/prospective.rs @@ -67,7 +67,7 @@ where _ = ticker.tick() => { let due_now = crate::memory::now_timestamp(); let mut guard = store.lock().await; - let _ = execute_due_reminders( + if let Err(error) = execute_due_reminders( &mut *guard, ExecuteDueRemindersRequest { tenant_id: tenant_id.clone(), @@ -79,7 +79,9 @@ where retry_strategy_id: None, trace_id: None, }, - ); + ) { + eprintln!("ProspectiveScheduler: failed to execute due reminders: {error}"); + } } } } diff --git a/src/runtime/retrieval.rs b/src/runtime/retrieval.rs index 75c4ccd..7dd86ac 100644 --- a/src/runtime/retrieval.rs +++ b/src/runtime/retrieval.rs @@ -1,4 +1,5 @@ use crate::{ + config::ScoringWeights, contracts::{CoreError, CoreResult}, memory::{estimate_tokens, MemoryRecord, PrivacyLevel}, runtime::intelligence::{ @@ -26,6 +27,7 @@ pub struct RetrievalRequest { pub lane: Option, pub policy_version: Option, pub trace_id: Option, + pub scoring_weights: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -116,6 +118,12 @@ impl RetrievalRequest { lane: None, policy_version: None, trace_id: None, + scoring_weights: Some(ScoringWeights { + semantic_similarity: 0.5, + recency: 0.2, + importance: 0.2, + access: 0.1, + }), } } } @@ -140,6 +148,12 @@ where if request.query_text.trim().is_empty() { return Err(CoreError::InvalidInput("query cannot be empty".to_string())); } + let scoring_weights = request.scoring_weights.clone().unwrap_or(ScoringWeights { + semantic_similarity: 0.5, + recency: 0.2, + importance: 0.2, + access: 0.1, + }); let trace_id = request.trace_id.clone().unwrap_or_else(|| { crate::memory::deterministic_id(&[ &request.tenant_id, @@ -240,7 +254,7 @@ where for (record, semantic_similarity) in &vector_items { merged.insert( record.id.clone(), - item_from_record(record, SourcePath::Vector, *semantic_similarity), + item_from_record(record, SourcePath::Vector, *semantic_similarity, &scoring_weights), ); } for memory_id in &graph_ids { @@ -253,6 +267,7 @@ where record, SourcePath::Graph, lexical_score(&record.content, &request.query_text), + &scoring_weights, ) }); } @@ -360,14 +375,26 @@ fn item_from_record( record: &MemoryRecord, source_path: SourcePath, semantic_similarity: f32, + weights: &ScoringWeights, ) -> RetrievedItem { let access = (record.access_count as f32 / 10.0).min(1.0); - let recency = 1.0; + // Calculate recency based on how recently the record was created + // Decay factor: 1.0 for recent, decreasing over time (half-life ~30 days) + let now = crate::memory::now_timestamp().parse::().unwrap_or(0); + let created = record.created_at.parse::().unwrap_or(0); + let age_seconds = now.saturating_sub(created); + let half_life_seconds = 30 * 24 * 60 * 60; // 30 days + let recency = if age_seconds == 0 { + 1.0 + } else { + (0.5_f32).powf(age_seconds as f32 / half_life_seconds as f32) + }; let score = retrieval_score( semantic_similarity, recency, record.importance_score, access, + weights, ); RetrievedItem { memory_id: record.id.clone(), diff --git a/src/scoring/mod.rs b/src/scoring/mod.rs index ecda4bc..9b4a80d 100644 --- a/src/scoring/mod.rs +++ b/src/scoring/mod.rs @@ -1,4 +1,5 @@ use crate::{ + config::ScoringWeights, contracts::{CoreError, CoreResult}, memory::MemoryRecord, }; @@ -58,6 +59,10 @@ pub fn retrieval_score( recency: f32, importance: f32, access: f32, + weights: &ScoringWeights, ) -> f32 { - (0.5 * semantic_similarity) + (0.2 * recency) + (0.2 * importance) + (0.1 * access) + (weights.semantic_similarity * semantic_similarity) + + (weights.recency * recency) + + (weights.importance * importance) + + (weights.access * access) } diff --git a/src/testkit/mod.rs b/src/testkit/mod.rs index 85eb963..10bd254 100644 --- a/src/testkit/mod.rs +++ b/src/testkit/mod.rs @@ -27,8 +27,14 @@ pub trait MemoryIndexStore { pub trait GraphStore { fn merge_node(&mut self, node: GraphNode) -> CoreResult<()>; fn merge_edge(&mut self, edge: GraphEdge) -> CoreResult<()>; - fn graph_memory_ids(&self, user_id: &str, query: &str, max_hops: u8) - -> CoreResult>; + fn graph_memory_ids( + &self, + tenant_id: &str, + user_id: &str, + query: &str, + max_hops: u8, + privacy_scope: &[PrivacyLevel], + ) -> CoreResult>; } pub trait AuditSink {