diff --git a/README.md b/README.md index 5830060..ec88268 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,7 @@ Then allow the MCP tools so Claude can use them without prompting each time. Add "mcp__recalld__forget_memory", "mcp__recalld__find_similar_memories", "mcp__recalld__create_namespace", + "mcp__recalld__namespace_stats", "mcp__recalld__list_memories" ] } @@ -172,7 +173,7 @@ See [docs/benchmark.md](docs/benchmark.md) for full methodology, per-category br ## Usage modes -**MCP server (stdio)** -- Runs as a Model Context Protocol server for AI tools like Claude Code. Exposes 9 tools: `store_memory`, `store_memories`, `recall_memories`, `get_memory`, `reinforce_memory`, `forget_memory`, `find_similar_memories`, `create_namespace`, `list_memories`. +**MCP server (stdio)** -- Runs as a Model Context Protocol server for AI tools like Claude Code. Exposes 10 tools: `store_memory`, `store_memories`, `recall_memories`, `get_memory`, `reinforce_memory`, `forget_memory`, `find_similar_memories`, `create_namespace`, `namespace_stats`, `list_memories`. ```sh recalld mcp diff --git a/docs/architecture.md b/docs/architecture.md index 9eaf466..5882318 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -44,7 +44,7 @@ graph TD subgraph STORE["Storage Engine"] META["meta.db"] - VEC["vectors.dat"] + VEC["<ns>/vectors.dat"] TXT["fulltext.dat"] EDGES["edges.db"] end @@ -90,7 +90,7 @@ Shutdown is the reverse: signal background tasks, drain in-flight requests (5s t ### Core storage files -All data lives in a single directory (default: `~/.recalld/data/`). An exclusive file lock (`recalld.lock`) prevents multi-process access. There are four primary data files (meta.db, vectors.dat, fulltext.dat, edges.db) plus the FTS5 index (fts.db) and the lock file. +All data lives in a single directory (default: `~/.recalld/data/`). An exclusive file lock (`recalld.lock`) prevents multi-process access. There are four primary data stores (meta.db, `/vectors.dat`, fulltext.dat, edges.db) plus the FTS5 index (fts.db) and the lock file. ``` ~/.recalld/data/ @@ -98,8 +98,8 @@ All data lives in a single directory (default: `~/.recalld/data/`). An exclusive meta.db # redb B+tree -- memory records + secondary indexes fulltext.dat # append-only CRC32 log -- full_text payloads edges.db # redb B+tree -- graph edge persistence - ns_default.dat # mmap'd vector file for "default" namespace - ns_work.dat # mmap'd vector file for "work" namespace (example) + default/vectors.dat # mmap'd vector file for "default" namespace + work/vectors.dat # mmap'd vector file for "work" namespace (example) fts.db # SQLite FTS5 full-text search index ``` @@ -122,7 +122,7 @@ Configured with a 16 MB read cache covering top B+tree levels. #### vectors.dat (per-namespace, mmap) -Each namespace gets its own memory-mapped vector file (`ns_.dat`). Vectors are stored in fixed-size slots for O(1) random access. +Each namespace gets its own subdirectory containing a memory-mapped `vectors.dat` file (e.g., `default/vectors.dat`, `work/vectors.dat`). Vectors are stored in fixed-size slots for O(1) random access. **File format:** @@ -174,10 +174,10 @@ Persists graph edges separately from memory records. Composite keys enable effic Both directions are stored for O(1) neighbor lookups in either direction. Orphaned edges (referencing deleted memories) are cleaned up on startup. -### Why four files +### Why four stores - **meta.db**: needs transactions, range scans, secondary indexes -> embedded B+tree (redb). -- **vectors.dat**: needs zero-copy reads of large float arrays, O(1) slot access -> mmap. +- **`/vectors.dat`**: needs zero-copy reads of large float arrays, O(1) slot access -> mmap (one file per namespace). - **fulltext.dat**: large variable-length blobs, append-only pattern -> log-structured file. - **edges.db**: could live in meta.db, but separated to avoid write amplification -- edges are written frequently (auto-linking) while metadata records are written infrequently after creation. diff --git a/docs/guide.md b/docs/guide.md index 565a35b..e821218 100644 --- a/docs/guide.md +++ b/docs/guide.md @@ -746,7 +746,7 @@ Shows decay forecast, at-risk memories, and storage breakdown. ## 6. MCP Tools Reference -recalld exposes 9 MCP tools. These are available to any MCP client (Claude Code, Cursor, etc.) when recalld is configured as an MCP server. +recalld exposes 10 MCP tools. These are available to any MCP client (Claude Code, Cursor, etc.) when recalld is configured as an MCP server. ### `store_memory` @@ -842,6 +842,14 @@ Create a new memory namespace. | `desiredRetention` | number | No | Target retention rate 0.0-1.0 (default: 0.9). | | `decayRateMultiplier` | number | No | Per-namespace decay rate multiplier. 1.0 = normal, 2.0 = 2x slower, 0.0 = disabled. Omit to inherit global setting. | +### `namespace_stats` + +Get statistics for a memory namespace including total memory count, phase breakdown (full/summary/ghost), permastore count, average strength, edge count, and vector storage size. + +| Parameter | Type | Required | Description | +|---|---|---|---| +| `namespace` | string | No | Namespace name to get stats for (default: `"default"`). | + ### `list_memories` List memories in a namespace with pagination and optional filters. Unlike `recall_memories`, this does not require a search query -- it returns memories sorted by creation date (newest first). diff --git a/docs/mcp.md b/docs/mcp.md index ec33e07..f2dd210 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -46,6 +46,7 @@ By default, Claude Code will prompt you for approval each time an MCP tool is ca "mcp__recalld__forget_memory", "mcp__recalld__find_similar_memories", "mcp__recalld__create_namespace", + "mcp__recalld__namespace_stats", "mcp__recalld__list_memories" ] } @@ -131,8 +132,7 @@ Response: "phase": "Full", "strength": 1.0, "stability": 3.7145, - "createdAt": 1719187200000, - "createdAtFormatted": "2025-06-24 00:00:00 UTC" + "createdAt": "2025-06-24 00:00:00 UTC" } ``` @@ -196,7 +196,7 @@ Response: "phase": "Full", "strength": 1.0, "stability": 3.7145, - "createdAt": 1719187200000 + "createdAt": "2025-06-24 00:00:00 UTC" }, { "index": 1, @@ -205,7 +205,7 @@ Response: "phase": "Full", "strength": 1.0, "stability": 3.7145, - "createdAt": 1719187200000 + "createdAt": "2025-06-24 00:00:00 UTC" } ], "total": 2, @@ -281,10 +281,8 @@ Response (compact=false): "topics": ["coding-style"], "phase": "Full", "strength": 0.95, - "createdAt": 1719187200000, - "createdAtFormatted": "2025-06-24 00:00:00 UTC", - "lastAccessedAt": 1719273600000, - "lastAccessedAtFormatted": "2025-06-25 00:00:00 UTC" + "createdAt": "2025-06-24 00:00:00 UTC", + "lastAccessedAt": "2025-06-25 00:00:00 UTC" } ], "count": 1, @@ -341,10 +339,8 @@ Response: "phase": "Full", "strength": 0.95, "stability": 12.4, - "createdAt": 1719187200000, - "createdAtFormatted": "2025-06-24 00:00:00 UTC", - "lastAccessedAt": 1719273600000, - "lastAccessedAtFormatted": "2025-06-25 00:00:00 UTC", + "createdAt": "2025-06-24 00:00:00 UTC", + "lastAccessedAt": "2025-06-25 00:00:00 UTC", "isPermastore": false, "edgeCount": 3 } @@ -470,8 +466,8 @@ Response: "tags": ["type/project", "tech/rust"], "phase": "Full", "strength": 0.88, - "createdAt": 1719100800000, - "lastAccessedAt": 1719187200000 + "createdAt": "2025-06-23 00:00:00 UTC", + "lastAccessedAt": "2025-06-24 00:00:00 UTC" } ], "count": 1 @@ -497,16 +493,19 @@ Response: "namespace": "default", "threshold": 0.9, "clusters": [ - [ - { - "id": "a1b2c3d4-...", - "summary": "User prefers snake_case in Rust" - }, - { - "id": "c3d4e5f6-...", - "summary": "User likes snake_case for Rust code" - } - ] + { + "memories": [ + { + "id": "a1b2c3d4-...", + "summary": "User prefers snake_case in Rust" + }, + { + "id": "c3d4e5f6-...", + "summary": "User likes snake_case for Rust code" + } + ], + "maxSimilarity": 0.92 + } ], "clusterCount": 1 } @@ -548,8 +547,47 @@ Response: "name": "work-project", "embeddingDim": 1536, "memoryCount": 0, - "createdAt": 1719187200000, - "createdAtFormatted": "2025-06-24 00:00:00 UTC" + "createdAt": "2025-06-24 00:00:00 UTC" +} +``` + +--- + +### namespace_stats + +Get statistics for a memory namespace including total memory count, phase breakdown (full/summary/ghost), permastore count, average strength, edge count, and vector storage size. Use this to check how many memories exist or monitor namespace health. + +#### Parameters + +| Parameter | Type | Required | Default | Description | +|-----------|------|----------|---------|-------------| +| `namespace` | string | no | `"default"` | Namespace name to get stats for | + +#### Example + +Request: + +```json +{ + "namespace": "default" +} +``` + +Response: + +```json +{ + "namespace": "default", + "totalMemories": 42, + "phases": { + "full": 30, + "summary": 8, + "ghost": 4 + }, + "permastoreCount": 3, + "averageStrength": 0.72, + "edgeCount": 156, + "vectorStorageBytes": 1048576 } ``` @@ -591,15 +629,12 @@ Response: { "id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890", "summary": "User prefers snake_case for Rust code and camelCase for TypeScript", - "namespace": "default", - "tags": ["type/user-profile", "tech/rust", "tech/typescript"], - "phase": "Full", - "strength": 0.95, - "createdAt": 1719187200000, - "createdAtFormatted": "2025-06-24 00:00:00 UTC" + "entities": ["TypeScript"], + "topics": ["coding-style"], + "createdAt": "2025-06-24 00:00:00 UTC" } ], - "count": 1, + "total": 1, "limit": 10, "offset": 0 } diff --git a/src/api/adapters.rs b/src/api/adapters.rs index 1c457ea..2bc02fd 100644 --- a/src/api/adapters.rs +++ b/src/api/adapters.rs @@ -11,12 +11,14 @@ use std::sync::Arc; use async_trait::async_trait; use crate::cache::CacheManager; +use crate::config::RecalldConfig; use crate::embedding::EmbeddingProvider; use crate::graph::SharedGraph; use crate::model::{ - AccessKind, CachedRecord, DecayPhase, DiskRecord, MemoryId, NamespaceConfig, NamespaceId, + AccessKind, CachedRecord, DecayPhase, DiskRecord, MemoryId, NamespaceConfig, NamespaceId, Tag, + parse_structured_tags, }; -use crate::search::FlatVectorIndex; +use crate::search::{EntityIndex, FlatVectorIndex, FtsIndex}; use crate::storage::RedbStorageEngine; // Import the StorageEngine trait so its methods are in scope. use super::state; @@ -31,6 +33,8 @@ pub struct SearchPipelineAdapter { query_engine: Arc, embedding: Arc, vector_index: Arc>, + fts_index: Arc>, + entity_index: Arc>, } impl SearchPipelineAdapter { @@ -39,11 +43,15 @@ impl SearchPipelineAdapter { query_engine: Arc, embedding: Arc, vector_index: Arc>, + fts_index: Arc>, + entity_index: Arc>, ) -> Self { Self { query_engine, embedding, vector_index, + fts_index, + entity_index, } } } @@ -67,43 +75,86 @@ impl state::SearchPipeline for SearchPipelineAdapter { state::QueryInput::Text(t) => Some(t), state::QueryInput::Vector(_) => None, }; + + // Build proper filter from query tags/entities (Issue 2) + let mut require_tags: Vec = query + .include_tags + .iter() + .filter_map(|t| Tag::new(t).ok()) + .collect(); + // Convert entities to entity/ tags (Issue 4) + for e in &query.entities { + if let Ok(tag) = Tag::new(&format!("entity/{}", e.to_lowercase())) { + require_tags.push(tag); + } + } + let exclude_tags: Vec = query + .exclude_tags + .iter() + .filter_map(|t| Tag::new(t).ok()) + .collect(); + let phases: Vec = query + .decay_phases + .as_ref() + .map(|ps| { + ps.iter() + .filter_map(|&p| match p { + 0 => Some(DecayPhase::Full), + 1 => Some(DecayPhase::Summary), + 2 => Some(DecayPhase::Ghost), + _ => None, + }) + .collect() + }) + .unwrap_or_default(); + let pipeline_query = crate::search::SearchQuery { text, fts_query: None, - namespace: String::new(), - filter: crate::search::PipelineSearchFilter::default(), + namespace: query.namespace_name.clone(), // Issue 1: pass actual namespace name + filter: crate::search::PipelineSearchFilter { + require_tags, + exclude_tags, + phases, + min_strength: query.min_score, // Issue 2: pass min_strength + ..Default::default() + }, limit: query.k, - min_score: query.min_score.unwrap_or(0.0), + min_score: 0.0, include_ghosts: false, query_mode: crate::search::QueryMode::default(), - graph_depth: query.graph_depth as u8, - time_range_start: None, - time_range_end: None, - entities: Vec::new(), + graph_depth: (query.graph_depth as u8).min(3), + time_range_start: query.time_range_start, // Issue 3: pass time range + time_range_end: query.time_range_end, // Issue 3: pass time range + entities: query.entities.clone(), // Issue 4: pass entities }; let response = self.query_engine.search(pipeline_query).await?; Ok(response .results .into_iter() - .map(|r| state::ResolvedSearchResult { - memory: CachedRecord { - id: r.memory_id, - namespace_id: NamespaceId::UNSET, - created_at: r.created_at, - last_accessed_at: 0, - phase: r.phase, - strength: r.retrievability, - decay_strength: r.effective_r, - stability: r.stability, - difficulty: 5.0, - is_permastore: r.is_permastore, - summary: r.summary.unwrap_or_default(), - tags: r.tags, - edge_count: r.edge_count, - vector_slot: 0, - entities: Vec::new(), - }, - score: r.score.unwrap_or(0.0), + .map(|r| { + // Issue 5: Parse entities from tags + let metadata = parse_structured_tags(&r.tags); + state::ResolvedSearchResult { + memory: CachedRecord { + id: r.memory_id, + namespace_id: query.namespace_id, // Issue 6: use actual namespace_id + created_at: r.created_at, + last_accessed_at: r.last_accessed_at, + phase: r.phase, + strength: r.retrievability, + decay_strength: r.effective_r, + stability: r.stability, + difficulty: 5.0, + is_permastore: r.is_permastore, + summary: r.summary.unwrap_or_default(), + tags: r.tags, + edge_count: r.edge_count, + vector_slot: 0, + entities: metadata.entities, + }, + score: r.score.unwrap_or(0.0), + } }) .collect()) } @@ -164,6 +215,49 @@ impl state::SearchPipeline for SearchPipelineAdapter { // A simple health check: try to embed a trivial string. self.embedding.embed("health check").await.is_ok() } + + async fn fts_add( + &self, + namespace_id: NamespaceId, + memory_id: MemoryId, + summary: &str, + full_text: Option<&str>, + tags: &[String], + ) { + let fts = self.fts_index.lock().await; + if let Err(e) = fts.add(namespace_id, memory_id, summary, full_text, tags) { + tracing::warn!( + memory_id = %memory_id, + %e, + "FTS5 indexing failed (non-fatal)" + ); + } + } + + async fn fts_remove(&self, memory_id: MemoryId) { + let fts = self.fts_index.lock().await; + if let Err(e) = fts.remove(memory_id) { + tracing::warn!( + memory_id = %memory_id, + %e, + "FTS5 removal failed (non-fatal)" + ); + } + } + + async fn entity_index_add(&self, memory_id: MemoryId, entities: &[String]) { + if !entities.is_empty() { + let mut idx = self.entity_index.write().await; + idx.add(memory_id, entities); + } + } + + async fn entity_index_remove(&self, memory_id: MemoryId, entities: &[String]) { + if !entities.is_empty() { + let mut idx = self.entity_index.write().await; + idx.remove(memory_id, entities); + } + } } // ═══════════════════════════════════════════════════════════════════════ @@ -174,22 +268,15 @@ impl state::SearchPipeline for SearchPipelineAdapter { pub struct StorageEngineAdapter { storage: Arc>, cache: Arc, - #[allow(dead_code)] - embedding: Arc, } impl StorageEngineAdapter { - /// Creates a new adapter wrapping the storage engine, cache, and embedding provider. + /// Creates a new adapter wrapping the storage engine and cache. pub fn new( storage: Arc>, cache: Arc, - embedding: Arc, ) -> Self { - Self { - storage, - cache, - embedding, - } + Self { storage, cache } } } @@ -222,7 +309,7 @@ impl state::StorageEngine for StorageEngineAdapter { phase: DecayPhase::Full, strength: 1.0, decay_strength: 1.0, - stability: initial_stability.unwrap_or(3.7), + stability: initial_stability.unwrap_or(3.7145), difficulty: 5.0, is_permastore: 0, vector_slot: 0, @@ -281,66 +368,45 @@ impl state::StorageEngine for StorageEngineAdapter { async fn delete_memory(&self, id: MemoryId) -> Result { let storage = self.storage.clone(); - let existing_record = tokio::task::spawn_blocking({ - let storage = storage.clone(); - move || { - let storage_r = storage.read().map_err(|e| { - crate::storage::StorageError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - format!("storage lock poisoned: {e}"), - )) - })?; - storage_r.get_record(id) - } - }) - .await - .map_err(|e| { - crate::storage::StorageError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - format!("blocking task join error: {e}"), - )) - })??; + // Perform read, tombstone check, tombstone, and free_vector_slot + // atomically under a single write lock to prevent TOCTOU races. + // Without this, two concurrent deletes could both read the record + // as non-tombstoned, both tombstone it, and both free the same + // vector slot -- corrupting the free list into a self-referential + // cycle. + let deleted = tokio::task::spawn_blocking(move || { + let mut storage_w = storage.write().map_err(|e| { + crate::storage::StorageError::Io(std::io::Error::new( + std::io::ErrorKind::Other, + format!("storage lock poisoned: {e}"), + )) + })?; - let Some(existing_record) = existing_record else { - return Ok(false); - }; + // Check if the record exists and is not already tombstoned. + let existing_record = match storage_w.get_record(id)? { + Some(r) => r, + None => return Ok::(false), + }; - if existing_record.phase == DecayPhase::Tombstone { - return Ok(false); - } + if existing_record.phase == DecayPhase::Tombstone { + return Ok(false); + } - // Tombstone the record and free the vector slot in a single blocking task. - tokio::task::spawn_blocking({ - let storage = storage.clone(); - move || { - { - let storage_r = storage.read().map_err(|e| { - crate::storage::StorageError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - format!("storage lock poisoned: {e}"), - )) - })?; - storage_r.tombstone_memory(id)?; - } - { - let mut storage_w = storage.write().map_err(|e| { - crate::storage::StorageError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - format!("storage lock poisoned: {e}"), - )) - })?; - let ns_id = NamespaceId::new(existing_record.namespace_id); - if let Err(e) = storage_w.free_vector_slot(ns_id, existing_record.vector_slot) { - tracing::warn!( - memory_id = %id, - vector_slot = existing_record.vector_slot, - %e, - "vector slot free failed (non-fatal)" - ); - } - } - Ok::<_, crate::storage::StorageError>(()) + // Tombstone the record. + storage_w.tombstone_memory(id)?; + + // Free the vector slot. + let ns_id = NamespaceId::new(existing_record.namespace_id); + if let Err(e) = storage_w.free_vector_slot(ns_id, existing_record.vector_slot) { + tracing::warn!( + memory_id = %id, + vector_slot = existing_record.vector_slot, + %e, + "vector slot free failed (non-fatal)" + ); } + + Ok(true) }) .await .map_err(|e| { @@ -350,8 +416,10 @@ impl state::StorageEngine for StorageEngineAdapter { )) })??; - self.cache.invalidate(id).await; - Ok(true) + if deleted { + self.cache.invalidate(id).await; + } + Ok(deleted) } async fn namespace_stats( @@ -465,6 +533,16 @@ impl state::StorageEngine for StorageEngineAdapter { { return false; } + if let Some(start) = filter.time_range_start { + if record.created_at < start { + return false; + } + } + if let Some(end) = filter.time_range_end { + if record.created_at > end { + return false; + } + } true }) .map(|(_id, disk_record)| CachedRecord::from(&disk_record)) @@ -515,6 +593,43 @@ impl state::StorageEngine for StorageEngineAdapter { })? } + async fn list_memories_filtered( + &self, + namespace_id: NamespaceId, + require_tags: &[Tag], + time_range_start: Option, + time_range_end: Option, + offset: usize, + limit: usize, + ) -> Result<(Vec<(MemoryId, DiskRecord)>, u64), crate::storage::StorageError> { + let storage = self.storage.clone(); + let tags_owned: Vec = require_tags.to_vec(); + + tokio::task::spawn_blocking(move || { + let storage_r = storage.read().map_err(|e| { + crate::storage::StorageError::Io(std::io::Error::new( + std::io::ErrorKind::Other, + format!("storage lock poisoned: {e}"), + )) + })?; + storage_r.meta_store().list_memories_filtered( + namespace_id, + &tags_owned, + time_range_start, + time_range_end, + offset, + limit, + ) + }) + .await + .map_err(|e| { + crate::storage::StorageError::Io(std::io::Error::new( + std::io::ErrorKind::Other, + format!("blocking task join error: {e}"), + )) + })? + } + async fn ping(&self) -> bool { let storage = self.storage.clone(); tokio::task::spawn_blocking(move || storage.read().is_ok()) @@ -585,10 +700,15 @@ impl state::StorageEngine for StorageEngineAdapter { })? } - fn storage_path(&self) -> PathBuf { + fn storage_path(&self) -> Result { tokio::task::block_in_place(|| { - let storage_r = self.storage.read().expect("storage lock not poisoned"); - storage_r.db_path().to_path_buf() + let storage_r = self.storage.read().map_err(|e| { + crate::storage::StorageError::Io(std::io::Error::new( + std::io::ErrorKind::Other, + format!("storage lock poisoned: {e}"), + )) + })?; + Ok(storage_r.db_path().to_path_buf()) }) } } @@ -659,12 +779,32 @@ impl state::RecordCache for RecordCacheAdapter { /// Adapts `SharedGraph` to the API `RelationshipGraph` trait. pub struct RelationshipGraphAdapter { graph: SharedGraph, + vector_index: Arc>, + entity_index: Arc>, + storage: Arc>, + cache: Arc, + config: Arc, } impl RelationshipGraphAdapter { - /// Creates a new adapter wrapping the shared relationship graph. - pub fn new(graph: SharedGraph) -> Self { - Self { graph } + /// Creates a new adapter wrapping the shared relationship graph and + /// subsystems needed for post-creation linking. + pub fn new( + graph: SharedGraph, + vector_index: Arc>, + entity_index: Arc>, + storage: Arc>, + cache: Arc, + config: Arc, + ) -> Self { + Self { + graph, + vector_index, + entity_index, + storage, + cache, + config, + } } } @@ -678,7 +818,8 @@ impl state::RelationshipGraph for RelationshipGraphAdapter { ) -> Result<(), crate::graph::GraphError> { use crate::model::EdgeType; let etype = match edge_type { - "parent_child" | "ParentChild" => EdgeType::ParentChild, + "parent" | "parent_child" | "ParentChild" => EdgeType::ParentChild, + "supersedes" | "Supersedes" => EdgeType::Supersedes, "associative" | "Associative" => EdgeType::Associative, "causal" | "Causal" => EdgeType::Causal, "contradicts" | "Contradicts" => EdgeType::Contradicts, @@ -686,14 +827,152 @@ impl state::RelationshipGraph for RelationshipGraphAdapter { }; let mut graph = self.graph.write().await; graph.add_edge(from, to, etype, 1.0, false)?; + + // Persist edge to storage. + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + let persisted = crate::storage::PersistedEdge { + source: from, + target: to, + edge_type: etype, + weight: 1.0, + auto_created: false, + created_at: now_ms, + }; + drop(graph); // Release graph lock before storage + + let storage = self.storage.clone(); + let _ = tokio::task::spawn_blocking(move || { + if let Ok(storage_r) = storage.read() { + let _ = storage_r.batch_add_edges(&[persisted]); + } + }) + .await; + Ok(()) } - async fn remove_all_edges(&self, id: MemoryId) -> Result<(), crate::graph::GraphError> { + async fn tombstone_node(&self, id: MemoryId) -> Result<(), crate::graph::GraphError> { let mut graph = self.graph.write().await; - graph.remove_node(id)?; + graph.update_node_state(id, DecayPhase::Tombstone, 0.0)?; Ok(()) } + + async fn add_node( + &self, + id: MemoryId, + namespace_id: NamespaceId, + phase: DecayPhase, + strength: f32, + vector_slot: u32, + ) -> Result<(), crate::graph::GraphError> { + let mut graph = self.graph.write().await; + let _ = graph.add_node(id, namespace_id, phase, strength, vector_slot); + Ok(()) + } + + async fn perform_post_creation_links( + &self, + memory_id: MemoryId, + namespace_id: NamespaceId, + embedding: &[f32], + tags: &[String], + entities: &[String], + created_at: i64, + ) { + if !self.config.graph.autolink_enabled { + return; + } + + let threshold = self.config.graph.auto_link_threshold as f32; + let max_links = self.config.graph.max_auto_links; + + // Autolink: discover and create edges to similar existing memories. + if let Err(e) = crate::graph::perform_autolink( + memory_id, + namespace_id, + embedding, + tags, + threshold, + max_links, + &self.vector_index, + &self.graph, + &self.storage, + &self.cache, + ) + .await + { + tracing::warn!( + memory_id = %memory_id, + %e, + "autolink failed (non-fatal)" + ); + } + + // Entity linking: connect memories sharing named entities. + if !entities.is_empty() { + let max_entity_links = self.config.graph.max_entity_links; + let _ = crate::graph::perform_entity_link( + memory_id, + namespace_id, + entities, + max_entity_links, + &self.entity_index, + &self.graph, + &self.storage, + &self.cache, + ) + .await; + } + + // Temporal linking: connect to recently-stored memories. + let temporal_window_ms = self.config.graph.temporal_window_ms; + let max_temporal_links = self.config.graph.max_temporal_links; + if temporal_window_ms > 0 { + let recent_memories: Vec<(MemoryId, i64)> = self + .cache + .iter() + .filter_map(|(mid, record)| { + if mid == memory_id { + return None; + } + if record.namespace_id != namespace_id { + return None; + } + let delta = (created_at - record.created_at).unsigned_abs(); + if delta <= temporal_window_ms { + Some((mid, record.created_at)) + } else { + None + } + }) + .collect(); + + if !recent_memories.is_empty() { + if let Err(e) = crate::graph::perform_temporal_link( + memory_id, + namespace_id, + created_at, + temporal_window_ms, + max_temporal_links, + &recent_memories, + &self.graph, + &self.storage, + &self.cache, + ) + .await + { + tracing::warn!( + memory_id = %memory_id, + %e, + "temporal link failed (non-fatal)" + ); + } + } + } + } } // ═══════════════════════════════════════════════════════════════════════ @@ -703,14 +982,26 @@ impl state::RelationshipGraph for RelationshipGraphAdapter { /// Adapts the decay subsystem to the API `FsrsEngine` trait. pub struct FsrsEngineAdapter { storage: Arc>, + cache: Arc, + graph: crate::graph::SharedGraph, + config: Arc, sweep_alive: bool, } impl FsrsEngineAdapter { - /// Creates a new adapter with the storage engine and sweep-thread liveness flag. - pub fn new(storage: Arc>, has_sweep: bool) -> Self { + /// Creates a new adapter with the storage engine, cache, and sweep-thread liveness flag. + pub fn new( + storage: Arc>, + cache: Arc, + graph: crate::graph::SharedGraph, + config: Arc, + has_sweep: bool, + ) -> Self { Self { storage, + cache, + graph, + config, sweep_alive: has_sweep, } } @@ -731,16 +1022,109 @@ impl state::FsrsEngine for FsrsEngineAdapter { async fn reinforce( &self, - _id: MemoryId, - _quality: u8, + id: MemoryId, + quality: u8, ) -> Result> { - // Reinforcement requires FSRS stability recalculation. - // Return a placeholder result reflecting current state. + // Step 1: Read the current record. + let record = { + let storage = self.storage.clone(); + tokio::task::spawn_blocking(move || { + let storage_r = storage + .read() + .map_err(|e| format!("storage lock poisoned: {e}"))?; + storage_r + .get_record(id) + .map_err(|e| e.to_string())? + .ok_or_else(|| format!("memory {id} not found")) + }) + .await + .map_err(|e| format!("blocking task join error: {e}"))?? + }; + + // Tombstoned memories cannot be reinforced. + if record.phase == DecayPhase::Tombstone { + return Err(format!("memory {id} has been deleted (tombstoned)").into()); + } + + let phase = record.phase; + + // Step 2: Look up the namespace's decay rate multiplier. + let ns_decay_multiplier = { + let storage = self.storage.clone(); + let ns_id = record.namespace_id; + tokio::task::spawn_blocking(move || { + let storage_r = storage + .read() + .map_err(|e| format!("storage lock poisoned: {e}"))?; + let multiplier = storage_r + .get_namespace(NamespaceId::new(ns_id)) + .ok() + .flatten() + .and_then(|ns| ns.decay_rate_multiplier) + .unwrap_or(1.0); + Ok::(multiplier) + }) + .await + .map_err(|e| format!("blocking task join error: {e}"))?? + }; + + // Step 3: Compute elapsed days since last access. + let now = chrono::Utc::now().timestamp_millis(); + let elapsed_millis = (now - record.last_accessed_at).max(0) as f64; + let elapsed_days = (elapsed_millis / 86_400_000.0) as f32; + + // Step 4: Use FSRS to compute new stability based on quality rating. + let decay_config = crate::decay::DecayConfig::default(); + let engine = crate::decay::FsrsEngine::new(&decay_config); + let new_stability = + engine.review_stability(record.stability, elapsed_days, quality, ns_decay_multiplier); + + // Step 5: Compute new retrievability (just reinforced = 1.0). + let new_strength = 1.0_f32; + let is_permastore = + record.is_permastore != 0 || new_stability >= decay_config.permastore_threshold; + + // Step 6: Persist the updated decay state and access event. + { + let storage = self.storage.clone(); + tokio::task::spawn_blocking(move || { + let storage_r = storage + .read() + .map_err(|e| format!("storage lock poisoned: {e}"))?; + storage_r + .update_decay_state( + id, + phase, + new_strength, + new_strength, + new_stability, + is_permastore, + ) + .map_err(|e| e.to_string())?; + storage_r + .update_access(id, now, crate::model::AccessKind::ManualReinforcement) + .map_err(|e| e.to_string()) + }) + .await + .map_err(|e| format!("blocking task join error: {e}"))??; + } + + // Step 7: Update cache. + if let Some(existing) = self.cache.get(id).await { + let mut updated = (*existing).clone(); + updated.stability = new_stability; + updated.strength = new_strength; + updated.decay_strength = new_strength; + updated.is_permastore = is_permastore; + updated.last_accessed_at = now; + self.cache.insert(id, updated).await; + } + Ok(state::ReinforceResult { - strength: 1.0, - stability: 3.7, - phase: DecayPhase::Full, - is_permastore: false, + strength: new_strength, + stability: new_stability, + phase, + is_permastore, }) } @@ -751,6 +1135,33 @@ impl state::FsrsEngine for FsrsEngineAdapter { fn last_sweep_time(&self) -> Option { None } + + async fn trigger_sweep( + &self, + as_of_millis: Option, + ) -> Result> { + use crate::decay::sweep::{DecaySweepRunner, SweepConfig}; + use crate::graph::ActivationConfig; + + let sweep_config = SweepConfig::default(); + let decay_config = crate::decay::DecayConfig::default(); + let activation_config = ActivationConfig::default(); + let global_multiplier = self.config.decay.decay_rate_multiplier as f64; + + let result = DecaySweepRunner::execute_sweep_at( + &sweep_config, + &decay_config, + &activation_config, + &self.storage, + &self.graph, + &self.cache, + global_multiplier, + as_of_millis, + ) + .await; + + Ok(result) + } } // ═══════════════════════════════════════════════════════════════════════ @@ -838,26 +1249,40 @@ impl state::NamespaceRegistry for NamespaceRegistryAdapter { async fn create( &self, name: &str, - embedding_dim: u32, + embedding_dim: Option, initial_stability: f32, desired_retention: f32, + decay_rate_multiplier: Option, ) -> Result> { let now = chrono::Utc::now().timestamp_millis(); - let ns_config = NamespaceConfig { - id: NamespaceId::UNSET, - name: name.to_string(), - embedding_dim, - initial_stability, - default_difficulty: 5.0, - phase_thresholds: crate::model::namespace::PhaseThresholds::default(), - permastore_threshold: 1500.0, - created_at: now, - desired_retention, - decay_rate_multiplier: None, - }; - let storage = self.storage.clone(); + let name_owned = name.to_string(); + tokio::task::spawn_blocking(move || { + // Resolve embedding_dim: use provided value, or inherit from + // the 'default' namespace, or fall back to 1536. + let dim = embedding_dim.unwrap_or_else(|| { + storage + .read() + .ok() + .and_then(|s| s.get_namespace_by_name("default").ok().flatten()) + .map(|ns| ns.embedding_dim) + .unwrap_or(1536) + }); + + let ns_config = NamespaceConfig { + id: NamespaceId::UNSET, + name: name_owned, + embedding_dim: dim, + initial_stability, + default_difficulty: 5.0, + phase_thresholds: crate::model::namespace::PhaseThresholds::default(), + permastore_threshold: 1500.0, + created_at: now, + desired_retention, + decay_rate_multiplier, + }; + let mut storage_w = storage.write().map_err(|e| { Box::new(std::io::Error::new( std::io::ErrorKind::Other, diff --git a/src/api/handlers.rs b/src/api/handlers.rs index 24648fe..1dc0d3a 100644 --- a/src/api/handlers.rs +++ b/src/api/handlers.rs @@ -19,6 +19,7 @@ use super::errors::AppError; use super::models::*; use super::state::{AppState, QueryInput, SearchQuery}; use crate::health::report as health_report_compute; +use crate::model::constants::NAMESPACE_NAME_MAX_BYTES; use crate::model::id::{MemoryId, NamespaceId}; use crate::model::memory::AccessKind; use crate::serialization::{ @@ -101,7 +102,28 @@ pub async fn create_memory( id: req.namespace.clone(), })?; - // --- Embedding --- + // --- Merge entities/topics/emotions into tags (Issue 8) --- + let mut merged_tags = req.tags.clone(); + for entity in &req.entities { + let tag = format!("entity/{}", entity.to_lowercase()); + if !merged_tags.contains(&tag) { + merged_tags.push(tag); + } + } + for topic in &req.topics { + let tag = format!("topic/{}", topic.to_lowercase()); + if !merged_tags.contains(&tag) { + merged_tags.push(tag); + } + } + for emotion in &req.emotions { + let tag = format!("emotion/{}", emotion.to_lowercase()); + if !merged_tags.contains(&tag) { + merged_tags.push(tag); + } + } + + // --- Embedding (Issue 9: embed summary + full_text + tags) --- let embedding = match req.embedding { Some(ref vec) => { if vec.len() != ns.embedding_dim as usize { @@ -117,33 +139,105 @@ pub async fn create_memory( } vec.clone() } - None => state.search.embed_text(&req.summary, ns.id).await?, + None => { + // Build embedding text from summary + full_text + tags (matches MCP) + let mut embed_text = match &req.full_text { + Some(ft) => format!("{}\n\n{}", req.summary, ft), + None => req.summary.clone(), + }; + if !merged_tags.is_empty() { + embed_text = format!("{} {}", embed_text, merged_tags.join(" ")); + } + state.search.embed_text(&embed_text, ns.id).await? + } }; // --- Persist --- + // Resolve initial stability: use caller-provided value, or fall back + // to the namespace's configured initial_stability (matching the MCP + // reference implementation). + let resolved_stability = req.initial_stability.unwrap_or(ns.initial_stability); let memory = state .storage .create_memory( ns.id, &req.summary, req.full_text.as_deref(), - &req.tags, + &merged_tags, &embedding, - req.initial_stability, + Some(resolved_stability), req.created_at, ) .await?; + // --- Cache + vector index --- + state.cache.insert(&memory).await; + state + .search + .index_memory(memory.id, &embedding, ns.id) + .await; + + // --- FTS5 indexing (Issue 10) --- + state + .search + .fts_add( + ns.id, + memory.id, + &req.summary, + req.full_text.as_deref(), + &merged_tags, + ) + .await; + + // --- Entity index (Issue 10) --- + state + .search + .entity_index_add(memory.id, &req.entities) + .await; + + // --- Graph node (Issue 11: must happen before edges) --- + let _ = state + .graph + .add_node( + memory.id, + ns.id, + crate::model::DecayPhase::Full, + 1.0, + memory.vector_slot, + ) + .await; + // --- Parent edge --- if let Some(parent_id) = req.parent_id { state.graph.add_edge(parent_id, memory.id, "parent").await?; } - // --- Cache + vector index --- - state.cache.insert(&memory).await; + // --- Supersedes edge (Issue 8) --- + if let Some(old_id) = req.supersedes { + if let Err(e) = state.graph.add_edge(memory.id, old_id, "supersedes").await { + tracing::warn!( + memory_id = %memory.id, + superseded = %old_id, + %e, + "supersedes edge failed (non-fatal)" + ); + } + } + + // --- Autolink, entity-link, temporal-link (Issue 12) --- + let created_at = req + .created_at + .unwrap_or_else(|| chrono::Utc::now().timestamp_millis()); state - .search - .index_memory(memory.id, &embedding, ns.id) + .graph + .perform_post_creation_links( + memory.id, + ns.id, + &embedding, + &merged_tags, + &req.entities, + created_at, + ) .await; let mut response = MemoryResponse::from_cached(&memory, ns.name.clone()); @@ -220,11 +314,34 @@ pub async fn delete_memory( let memory_id = parse_memory_id(&id)?; + // Read record first for entity index cleanup (Issue 14) + let existing_record = state.storage.get_record(memory_id).await; + let existed = state.storage.delete_memory(memory_id).await?; if existed { state.cache.remove(memory_id).await; state.search.remove_from_index(memory_id).await; - state.graph.remove_all_edges(memory_id).await?; + + // Issue 14: Clean up FTS5 index + state.search.fts_remove(memory_id).await; + + // Issue 14: Clean up entity index + if let Some(ref record) = existing_record { + let metadata = crate::model::parse_structured_tags(&record.tags); + state + .search + .entity_index_remove(memory_id, &metadata.entities) + .await; + } + + // Issue 13: Tombstone the graph node instead of removing it + if let Err(e) = state.graph.tombstone_node(memory_id).await { + tracing::warn!( + memory_id = %memory_id, + %e, + "graph tombstone failed (non-fatal)" + ); + } } let took = start.elapsed().as_micros() as u64; @@ -356,25 +473,87 @@ pub async fn list_memories( None }; - // Convert phase string to DecayPhase enum + // Issue 19: Use list_memories_filtered() for server-side pagination when namespace is given. + // Build require_tags from tags + entities (Issue 21). + let mut require_tags: Vec = params + .tags + .iter() + .filter_map(|t| crate::model::Tag::new(t).ok()) + .collect(); + for entity in ¶ms.entities { + if let Ok(tag) = crate::model::Tag::new(&format!("entity/{}", entity.to_lowercase())) { + require_tags.push(tag); + } + } + + if let Some(ns_id) = namespace_id { + // Server-side filtered pagination (Issue 19) + let (page, total) = state + .storage + .list_memories_filtered( + ns_id, + &require_tags, + params.time_range_start, + params.time_range_end, + offset, + limit, + ) + .await?; + + let mut memories: Vec = Vec::with_capacity(page.len()); + for (mid, record) in page { + let ns_name = state + .namespaces + .name_for(NamespaceId::new(record.namespace_id)) + .unwrap_or_else(|| "unknown".to_string()); + let cached = crate::model::CachedRecord::from(&record); + let mut mem = MemoryResponse::from_cached(&cached, ns_name); + mem.full_text = state.storage.get_full_text(mid).await.unwrap_or(None); + memories.push(mem); + } + + let has_more = (offset + memories.len()) < total as usize; + + let response = ListMemoriesResponse { + memories, + total, + limit: limit as u32, + offset: offset as u32, + has_more, + }; + + let took = start.elapsed().as_micros() as u64; + + return Ok(Json(ApiResponse { + data: response, + took_us: Some(took), + })); + } + + // Fallback: no namespace specified, use in-memory filtering let phase_filter = params.phase.as_ref().map(|p| match p.as_str() { "full" => crate::model::DecayPhase::Full, "summary" => crate::model::DecayPhase::Summary, "ghost" => crate::model::DecayPhase::Ghost, - _ => crate::model::DecayPhase::Full, // unreachable due to validation above + _ => crate::model::DecayPhase::Full, }); - // Build filter struct let filter = ListFilter { namespace_id, phase: phase_filter, - tags: params.tags, + tags: { + let mut all_tags = params.tags; + for entity in ¶ms.entities { + all_tags.push(format!("entity/{}", entity.to_lowercase())); + } + all_tags + }, + time_range_start: params.time_range_start, + time_range_end: params.time_range_end, }; - // Get filtered records from storage let mut records = state.storage.list_memories(&filter).await?; - // Sort in memory match (sort_field, order) { ("created", "asc") => records.sort_by_key(|r| r.created_at), ("created", "desc") => { @@ -404,15 +583,13 @@ pub async fn list_memories( .partial_cmp(&a.stability) .unwrap_or(std::cmp::Ordering::Equal) }), - _ => {} // unreachable + _ => {} } let total = records.len() as u64; - // Apply pagination let page_records: Vec<_> = records.into_iter().skip(offset).take(limit).collect(); - // Convert to MemoryResponse objects, loading full_text from storage let mut memories: Vec = Vec::with_capacity(page_records.len()); for record in page_records { let ns_name = state @@ -515,7 +692,21 @@ pub async fn search_memories( } } - // Build internal query + // Merge entities/topics/emotions into tag filters (matching create_memory behaviour) + let mut include_tags = req.tags.clone(); + for topic in &req.topics { + let tag = format!("topic/{}", topic.to_lowercase()); + if !include_tags.contains(&tag) { + include_tags.push(tag); + } + } + for emotion in &req.emotions { + let tag = format!("emotion/{}", emotion.to_lowercase()); + if !include_tags.contains(&tag) { + include_tags.push(tag); + } + } + let query_input = if let Some(ref text) = req.query { Some(QueryInput::Text(text.clone())) } else { @@ -527,13 +718,17 @@ pub async fn search_memories( let search_query = SearchQuery { query: query_input.unwrap_or(QueryInput::Text(String::new())), namespace_id: ns.id, + namespace_name: ns.name.clone(), k: limit, - include_tags: req.tags.clone(), + include_tags, exclude_tags: vec![], decay_phases: None, min_score: req.min_strength, graph_depth: req.depth as usize, apply_rif: true, + time_range_start: req.time_range_start, + time_range_end: req.time_range_end, + entities: req.entities.clone(), }; let results = state.search.search(search_query).await?; @@ -615,11 +810,17 @@ pub async fn find_similar( let limit = req.limit.min(100) as usize; + let ns_name = state + .namespaces + .name_for(source.namespace_id) + .unwrap_or_else(|| "default".to_string()); + // Search using source embedding, requesting limit+1 to account // for the source memory appearing in its own results. let search_query = SearchQuery { query: QueryInput::Vector(source_embedding), namespace_id: source.namespace_id, + namespace_name: ns_name, k: limit + 1, include_tags: vec![], exclude_tags: vec![], @@ -627,6 +828,9 @@ pub async fn find_similar( min_score: req.min_score, graph_depth: 0, apply_rif: false, + time_range_start: None, + time_range_end: None, + entities: vec![], }; let results = state.search.search(search_query).await?; @@ -721,9 +925,10 @@ pub async fn create_namespace( let start = Instant::now(); // Validate name format - if req.name.is_empty() || req.name.len() > 64 { + if req.name.is_empty() || req.name.len() > NAMESPACE_NAME_MAX_BYTES { return Err(AppError::BadRequest { - message: "namespace name must be 1-64 characters".into(), + message: format!("namespace name must be 1-{NAMESPACE_NAME_MAX_BYTES} characters") + .into(), field: Some("name".into()), }); } @@ -740,12 +945,14 @@ pub async fn create_namespace( }); } - // Validate embedding dimensions - if req.embedding_dim == 0 || req.embedding_dim > 8192 { - return Err(AppError::BadRequest { - message: "embedding_dim must be between 1 and 8192".into(), - field: Some("embeddingDim".into()), - }); + // Validate embedding dimensions if provided + if let Some(dim) = req.embedding_dim { + if dim == 0 || dim > 8192 { + return Err(AppError::BadRequest { + message: "embedding_dim must be between 1 and 8192".into(), + field: Some("embeddingDim".into()), + }); + } } // Create namespace (registry checks for duplicates) @@ -756,6 +963,7 @@ pub async fn create_namespace( req.embedding_dim, req.initial_stability, req.desired_retention, + req.decay_rate_multiplier, ) .await .map_err(|_e| AppError::Conflict { @@ -1027,6 +1235,338 @@ pub async fn health_report( })) } +/// POST /memories/batch -- batch create memories (Issue 16). +/// +/// Creates multiple memories in sequence. Each memory goes through the +/// same validation and creation pipeline as the single-create endpoint. +pub async fn batch_store( + State(state): State, + Json(req): Json, +) -> Result<(StatusCode, Json>), AppError> { + let start = Instant::now(); + + if req.memories.is_empty() { + return Err(AppError::BadRequest { + message: "memories array must not be empty".into(), + field: Some("memories".into()), + }); + } + + if req.memories.len() > 100 { + return Err(AppError::BadRequest { + message: "batch size exceeds 100".into(), + field: Some("memories".into()), + }); + } + + let mut created = Vec::with_capacity(req.memories.len()); + + for mem_req in req.memories { + // Validate + if mem_req.summary.is_empty() { + continue; + } + + let ns = match state.namespaces.resolve(&mem_req.namespace) { + Some(ns) => ns, + None => continue, + }; + + // Merge entities/topics/emotions into tags + let mut merged_tags = mem_req.tags.clone(); + for entity in &mem_req.entities { + let tag = format!("entity/{}", entity.to_lowercase()); + if !merged_tags.contains(&tag) { + merged_tags.push(tag); + } + } + for topic in &mem_req.topics { + let tag = format!("topic/{}", topic.to_lowercase()); + if !merged_tags.contains(&tag) { + merged_tags.push(tag); + } + } + for emotion in &mem_req.emotions { + let tag = format!("emotion/{}", emotion.to_lowercase()); + if !merged_tags.contains(&tag) { + merged_tags.push(tag); + } + } + + // Embedding + let embedding = match mem_req.embedding { + Some(ref vec) => { + if vec.len() != ns.embedding_dim as usize { + continue; + } + vec.clone() + } + None => { + let mut embed_text = match &mem_req.full_text { + Some(ft) => format!("{}\n\n{}", mem_req.summary, ft), + None => mem_req.summary.clone(), + }; + if !merged_tags.is_empty() { + embed_text = format!("{} {}", embed_text, merged_tags.join(" ")); + } + match state.search.embed_text(&embed_text, ns.id).await { + Ok(emb) => emb, + Err(_) => continue, + } + } + }; + + // Persist — resolve initial stability from namespace config (matching MCP) + let resolved_stability = mem_req.initial_stability.unwrap_or(ns.initial_stability); + let memory = match state + .storage + .create_memory( + ns.id, + &mem_req.summary, + mem_req.full_text.as_deref(), + &merged_tags, + &embedding, + Some(resolved_stability), + mem_req.created_at, + ) + .await + { + Ok(m) => m, + Err(_) => continue, + }; + + // Cache + index + state.cache.insert(&memory).await; + state + .search + .index_memory(memory.id, &embedding, ns.id) + .await; + + // FTS + state + .search + .fts_add( + ns.id, + memory.id, + &mem_req.summary, + mem_req.full_text.as_deref(), + &merged_tags, + ) + .await; + + // Entity index + state + .search + .entity_index_add(memory.id, &mem_req.entities) + .await; + + // Graph node + let _ = state + .graph + .add_node( + memory.id, + ns.id, + crate::model::DecayPhase::Full, + 1.0, + memory.vector_slot, + ) + .await; + + // Supersedes edge + if let Some(old_id) = mem_req.supersedes { + let _ = state.graph.add_edge(memory.id, old_id, "supersedes").await; + } + + // Post-creation links + let created_at = mem_req + .created_at + .unwrap_or_else(|| chrono::Utc::now().timestamp_millis()); + state + .graph + .perform_post_creation_links( + memory.id, + ns.id, + &embedding, + &merged_tags, + &mem_req.entities, + created_at, + ) + .await; + + created.push(BatchStoreResult { + id: memory.id, + namespace: ns.name.clone(), + }); + } + + let total = created.len() as u64; + let took = start.elapsed().as_micros() as u64; + + Ok(( + StatusCode::CREATED, + Json(ApiResponse { + data: BatchStoreResponse { created, total }, + took_us: Some(took), + }), + )) +} + +/// POST /namespaces/:name/duplicates -- scan for duplicate memories (Issue 16). +/// +/// Uses vector similarity to find clusters of near-duplicate memories +/// within a namespace. +pub async fn scan_duplicates( + State(state): State, + Path(name): Path, + Json(req): Json, +) -> Result>, AppError> { + let start = Instant::now(); + + let ns = state + .namespaces + .resolve(&name) + .ok_or_else(|| AppError::NotFound { + resource: "namespace", + id: name.clone(), + })?; + + let threshold = req.threshold.clamp(0.0, 1.0); + let max_memories = req.max_memories.min(10_000); + + // Get memory IDs in this namespace + let filter = ListFilter { + namespace_id: Some(ns.id), + phase: None, + tags: vec![], + time_range_start: None, + time_range_end: None, + }; + let records = state.storage.list_memories(&filter).await?; + let memory_ids: Vec = records.iter().take(max_memories).map(|r| r.id).collect(); + + if memory_ids.len() < 2 { + let took = start.elapsed().as_micros() as u64; + return Ok(Json(ApiResponse { + data: ScanDuplicatesResponse { + clusters: vec![], + total: 0, + }, + took_us: Some(took), + })); + } + + // Load embeddings + let embeddings = state.search.get_embeddings_batch(&memory_ids); + let id_embeddings: Vec<(MemoryId, Vec)> = memory_ids + .into_iter() + .filter_map(|mid| embeddings.get(&mid).map(|emb| (mid, emb.clone()))) + .collect(); + + if id_embeddings.len() < 2 { + let took = start.elapsed().as_micros() as u64; + return Ok(Json(ApiResponse { + data: ScanDuplicatesResponse { + clusters: vec![], + total: 0, + }, + took_us: Some(took), + })); + } + + // Union-Find based clustering + let n = id_embeddings.len(); + let mut parent: Vec = (0..n).collect(); + + fn find(parent: &mut [usize], mut i: usize) -> usize { + while parent[i] != i { + parent[i] = parent[parent[i]]; + i = parent[i]; + } + i + } + + fn union(parent: &mut [usize], a: usize, b: usize) { + let ra = find(parent, a); + let rb = find(parent, b); + if ra != rb { + parent[rb] = ra; + } + } + + let mut max_sim: std::collections::HashMap = std::collections::HashMap::new(); + + // Compare each memory's embedding against its nearest neighbors + for (i, (mid_i, _emb_i)) in id_embeddings.iter().enumerate() { + // Find similar via the search pipeline's get_embedding + manual cosine + if let Some(emb) = state.search.get_embedding(*mid_i) { + // Use pairwise comparison within the loaded set + for (j, (_mid_j, emb_j)) in id_embeddings.iter().enumerate() { + if i >= j { + continue; + } + let score = crate::search::dot_product_simd(&emb, emb_j); + if score >= threshold { + union(&mut parent, i, j); + let root = find(&mut parent, i); + let entry = max_sim.entry(root).or_insert(0.0); + if score > *entry { + *entry = score; + } + } + } + } + } + + // Group into clusters + let mut clusters_map: std::collections::HashMap> = + std::collections::HashMap::new(); + for i in 0..n { + let root = find(&mut parent, i); + clusters_map.entry(root).or_default().push(i); + } + + let mut clusters: Vec = Vec::new(); + for (root, members) in &clusters_map { + if members.len() < 2 { + continue; + } + let entries: Vec = members + .iter() + .map(|&idx| { + let mid = id_embeddings[idx].0; + let summary = records + .iter() + .find(|r| r.id == mid) + .map(|r| r.summary.clone()) + .unwrap_or_default(); + DuplicateEntry { + id: mid.to_string(), + summary, + } + }) + .collect(); + let max_similarity = max_sim.get(root).copied().unwrap_or(0.0); + clusters.push(DuplicateCluster { + memories: entries, + max_similarity, + }); + } + + clusters.sort_by(|a, b| { + b.max_similarity + .partial_cmp(&a.max_similarity) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + let total = clusters.len() as u64; + let took = start.elapsed().as_micros() as u64; + + Ok(Json(ApiResponse { + data: ScanDuplicatesResponse { clusters, total }, + took_us: Some(took), + })) +} + /// GET /metrics -- Prometheus exposition format. /// /// Returns `text/plain` with metrics lines. Not JSON. @@ -1052,3 +1592,34 @@ pub async fn metrics( output, )) } + +/// POST /decay/sweep -- trigger an immediate decay sweep. +/// +/// Accepts an optional `as_of` field (milliseconds since epoch) to run +/// the sweep as if the current time were the given timestamp. This is +/// used by the replay driver to simulate time-appropriate decay during +/// historical seeding. +pub async fn trigger_decay_sweep( + State(state): State, + Json(req): Json, +) -> Result>, AppError> { + let result = state + .decay + .trigger_sweep(req.as_of) + .await + .map_err(|e| AppError::Internal { + source: format!("sweep failed: {e}").into(), + })?; + + Ok(Json(ApiResponse { + data: TriggerSweepResponse { + memories_scanned: result.memories_scanned, + full_to_summary: result.full_to_summary, + summary_to_ghost: result.summary_to_ghost, + deletions: result.deletions, + saved_by_connection_bonus: result.saved_by_connection_bonus, + duration_ms: result.duration.as_millis() as u64, + }, + took_us: Some(result.duration.as_micros() as u64), + })) +} diff --git a/src/api/mod.rs b/src/api/mod.rs index 75a615f..bed3e1e 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -33,7 +33,7 @@ pub use state::{ // Re-export API-layer supporting types. pub use state::{ NamespaceListInfo, NamespaceStats, QueryInput, ReinforceResult, ResolvedSearchResult, - SearchFilter, SearchQuery, + SearchQuery, }; use std::net::SocketAddr; diff --git a/src/api/models.rs b/src/api/models.rs index 807d1e2..3549e39 100644 --- a/src/api/models.rs +++ b/src/api/models.rs @@ -48,6 +48,24 @@ pub struct CreateMemoryApiRequest { #[serde(skip_serializing_if = "Option::is_none", default)] pub initial_stability: Option, + /// Named entities (e.g., person/place/org names). Merged into tags + /// as `entity/` entries. + #[serde(default)] + pub entities: Vec, + + /// Topics. Merged into tags as `topic/` entries. + #[serde(default)] + pub topics: Vec, + + /// Emotions. Merged into tags as `emotion/` entries. + #[serde(default)] + pub emotions: Vec, + + /// Optional ID of a memory this one supersedes. Creates a + /// supersedes edge from the new memory to the old memory. + #[serde(skip_serializing_if = "Option::is_none", default)] + pub supersedes: Option, + /// Optional parent memory ID. If provided, creates a parent->child /// edge from `parent_id` to the new memory. #[serde(skip_serializing_if = "Option::is_none", default)] @@ -132,11 +150,6 @@ pub struct FindSimilarRequest { #[serde(default)] pub min_score: Option, - /// Whether to search only within the same namespace as the source - /// memory. Default: `true`. - #[serde(default = "default_same_namespace")] - pub same_namespace: bool, - /// Whether to include embedding vectors in results. Default: `false`. #[serde(default)] pub include_embeddings: bool, @@ -147,11 +160,6 @@ fn default_similar_limit() -> u32 { 10 } -/// Default for `same_namespace` filter. -fn default_same_namespace() -> bool { - true -} - // ═══════════════════════════════════════════════════════════════════════ // Namespace Endpoints // ═══════════════════════════════════════════════════════════════════════ @@ -235,6 +243,18 @@ pub struct ListMemoriesQuery { #[serde(default, deserialize_with = "deserialize_comma_separated")] pub tags: Vec, + /// Filter by entities (comma-separated, AND logic). Converted to entity/ tags. + #[serde(default, deserialize_with = "deserialize_comma_separated")] + pub entities: Vec, + + /// Optional inclusive lower bound for temporal filtering (millis since epoch). + #[serde(default)] + pub time_range_start: Option, + + /// Optional inclusive upper bound for temporal filtering (millis since epoch). + #[serde(default)] + pub time_range_end: Option, + /// Sort field: "created", "accessed", "strength", "stability". #[serde(default)] pub sort: Option, @@ -292,6 +312,10 @@ pub struct ListFilter { pub phase: Option, /// Tag filter (AND logic: memory must have ALL tags). pub tags: Vec, + /// Optional inclusive lower bound for temporal filtering (millis since epoch). + pub time_range_start: Option, + /// Optional inclusive upper bound for temporal filtering (millis since epoch). + pub time_range_end: Option, } // ═══════════════════════════════════════════════════════════════════════ @@ -465,7 +489,7 @@ pub struct AgeDistribution { } /// Storage breakdown by file. -#[derive(Debug, Clone, Serialize)] +#[derive(Debug, Clone, Default, Serialize)] #[serde(rename_all = "camelCase")] pub struct StorageBreakdown { /// Total size of data directory in bytes. @@ -509,3 +533,111 @@ pub struct TagCount { /// Number of memories with this tag. pub count: u64, } + +// ═══════════════════════════════════════════════════════════════════════ +// Batch Store Endpoint (Issue 16) +// ═══════════════════════════════════════════════════════════════════════ + +/// POST /memories/batch -- batch create memories. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchStoreRequest { + /// List of memories to create. + pub memories: Vec, +} + +/// POST /memories/batch -- response. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchStoreResponse { + /// Created memory IDs. + pub created: Vec, + /// Total number successfully created. + pub total: u64, +} + +/// Result for a single batch store item. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchStoreResult { + /// Created memory ID. + pub id: MemoryId, + /// Namespace name. + pub namespace: String, +} + +// ═══════════════════════════════════════════════════════════════════════ +// Scan Duplicates Endpoint (Issue 16) +// ═══════════════════════════════════════════════════════════════════════ + +/// POST /namespaces/:name/duplicates -- scan for duplicate memories. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ScanDuplicatesRequest { + /// Similarity threshold (0.0 to 1.0). Default: 0.85. + #[serde(default = "default_duplicate_threshold")] + pub threshold: f32, + + /// Maximum memories to scan. Default: 500. + #[serde(default = "default_max_scan")] + pub max_memories: usize, +} + +fn default_duplicate_threshold() -> f32 { + 0.85 +} + +fn default_max_scan() -> usize { + 500 +} + +/// POST /namespaces/:name/duplicates -- response. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ScanDuplicatesResponse { + /// Clusters of similar memories. + pub clusters: Vec, + /// Total clusters found. + pub total: u64, +} + +/// A cluster of duplicate/near-duplicate memories. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DuplicateCluster { + /// Memories in this cluster. + pub memories: Vec, + /// Maximum similarity score within the cluster. + pub max_similarity: f32, +} + +/// A single memory entry in a duplicate cluster. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DuplicateEntry { + /// Memory ID. + pub id: String, + /// Memory summary. + pub summary: String, +} + +// ═══════════════════════════════════════════════════════════════════════ +// Decay sweep +// ═══════════════════════════════════════════════════════════════════════ + +#[derive(Debug, serde::Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TriggerSweepRequest { + pub as_of: Option, +} + +#[derive(Debug, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct TriggerSweepResponse { + pub memories_scanned: u64, + pub full_to_summary: u64, + pub summary_to_ghost: u64, + pub deletions: u64, + pub saved_by_connection_bonus: u64, + pub duration_ms: u64, +} diff --git a/src/api/routes.rs b/src/api/routes.rs index fa4e5ab..1c061e4 100644 --- a/src/api/routes.rs +++ b/src/api/routes.rs @@ -50,6 +50,7 @@ pub fn router(state: AppState, config: &ApiConfig) -> Router { let memory_routes = Router::new() .route("/", get(handlers::list_memories)) .route("/", post(handlers::create_memory)) + .route("/batch", post(handlers::batch_store)) .route("/{id}", get(handlers::get_memory)) .route("/{id}", delete(handlers::delete_memory)) .route("/{id}/reinforce", post(handlers::reinforce_memory)); @@ -61,12 +62,14 @@ pub fn router(state: AppState, config: &ApiConfig) -> Router { let namespace_routes = Router::new() .route("/", get(handlers::list_namespaces)) .route("/", post(handlers::create_namespace)) - .route("/{id}/stats", get(handlers::namespace_stats)); + .route("/{id}/stats", get(handlers::namespace_stats)) + .route("/{name}/duplicates", post(handlers::scan_duplicates)); let ops_routes = Router::new() .route("/health", get(handlers::health_check)) .route("/health/report", get(handlers::health_report)) - .route("/metrics", get(handlers::metrics)); + .route("/metrics", get(handlers::metrics)) + .route("/decay/sweep", post(handlers::trigger_decay_sweep)); let cors = if config.cors_permissive { CorsLayer::permissive() diff --git a/src/api/state.rs b/src/api/state.rs index 257c7be..f0fdc2c 100644 --- a/src/api/state.rs +++ b/src/api/state.rs @@ -62,6 +62,25 @@ pub trait SearchPipeline: Send + Sync { /// Check whether the embedding provider is healthy. async fn embedding_provider_healthy(&self) -> bool; + + /// Add a memory to the FTS5 full-text search index. + async fn fts_add( + &self, + namespace_id: NamespaceId, + memory_id: MemoryId, + summary: &str, + full_text: Option<&str>, + tags: &[String], + ); + + /// Remove a memory from the FTS5 full-text search index. + async fn fts_remove(&self, memory_id: MemoryId); + + /// Add a memory to the entity index for entity-based linking. + async fn entity_index_add(&self, memory_id: MemoryId, entities: &[String]); + + /// Remove a memory from the entity index. + async fn entity_index_remove(&self, memory_id: MemoryId, entities: &[String]); } /// Persistent storage engine (meta.db, fulltext.dat, vectors.dat, edges.db). @@ -102,6 +121,20 @@ pub trait StorageEngine: Send + Sync { filter: &crate::api::models::ListFilter, ) -> Result, StorageError>; + /// List memories with server-side filtering and pagination. + /// + /// Returns `(matching_records, total_count)` where the records + /// are the paginated slice and the count is total before pagination. + async fn list_memories_filtered( + &self, + namespace_id: NamespaceId, + require_tags: &[crate::model::tag::Tag], + time_range_start: Option, + time_range_end: Option, + offset: usize, + limit: usize, + ) -> Result<(Vec<(MemoryId, DiskRecord)>, u64), StorageError>; + /// Health probe: returns `true` if the storage engine can read/write. async fn ping(&self) -> bool; @@ -123,7 +156,7 @@ pub trait StorageEngine: Send + Sync { async fn list_tags(&self) -> Result, StorageError>; /// Return the database directory path for file-size computation. - fn storage_path(&self) -> PathBuf; + fn storage_path(&self) -> Result; } /// RAM record cache (moka-backed). @@ -164,8 +197,32 @@ pub trait RelationshipGraph: Send + Sync { edge_type: &str, ) -> Result<(), GraphError>; - /// Remove all edges (both directions) for a memory. - async fn remove_all_edges(&self, id: MemoryId) -> Result<(), GraphError>; + /// Tombstone a memory's graph node (set phase to Tombstone, strength to 0). + /// Unlike full removal, this preserves the node and edges for graph traversal. + async fn tombstone_node(&self, id: MemoryId) -> Result<(), GraphError>; + + /// Add a node to the relationship graph for a newly created memory. + async fn add_node( + &self, + id: MemoryId, + namespace_id: NamespaceId, + phase: crate::model::decay::DecayPhase, + strength: f32, + vector_slot: u32, + ) -> Result<(), GraphError>; + + /// Perform autolink, entity-link, and temporal-link for a newly created memory. + /// This is a composite operation that discovers and creates edges to similar, + /// entity-related, and temporally-adjacent memories. + async fn perform_post_creation_links( + &self, + memory_id: MemoryId, + namespace_id: NamespaceId, + embedding: &[f32], + tags: &[String], + entities: &[String], + created_at: i64, + ); } /// FSRS decay engine for reinforcement and strength recalculation. @@ -186,6 +243,12 @@ pub trait FsrsEngine: Send + Sync { /// Return the time of the last completed decay sweep, if any. fn last_sweep_time(&self) -> Option; + + /// Trigger an immediate decay sweep, optionally at a simulated time. + async fn trigger_sweep( + &self, + as_of_millis: Option, + ) -> Result>; } /// Namespace registry: name <-> id mapping, config lookup. @@ -207,9 +270,10 @@ pub trait NamespaceRegistry: Send + Sync { async fn create( &self, name: &str, - embedding_dim: u32, + embedding_dim: Option, initial_stability: f32, desired_retention: f32, + decay_rate_multiplier: Option, ) -> Result>; } @@ -238,6 +302,8 @@ pub struct SearchQuery { pub query: QueryInput, /// Target namespace for the search. pub namespace_id: NamespaceId, + /// Target namespace name (passed to pipeline for namespace-scoped search). + pub namespace_name: String, /// Maximum number of results to return. pub k: usize, /// Tag inclusion filter (AND semantics). @@ -252,6 +318,12 @@ pub struct SearchQuery { pub graph_depth: usize, /// Whether to apply retrieval-induced forgetting. pub apply_rif: bool, + /// Optional inclusive lower bound for temporal filtering (millis since epoch). + pub time_range_start: Option, + /// Optional inclusive upper bound for temporal filtering (millis since epoch). + pub time_range_end: Option, + /// Named entities for entity overlap scoring. + pub entities: Vec, } /// Search query input: either natural-language text or a pre-computed @@ -264,15 +336,6 @@ pub enum QueryInput { Vector(Vec), } -/// Search filter parameters for the similar-memories endpoint. -#[derive(Debug, Clone, Default)] -pub struct SearchFilter { - /// Optional namespace restriction. - pub namespace_id: Option, - /// Minimum score threshold. - pub min_score: Option, -} - /// A single resolved search result with its similarity score. #[derive(Debug, Clone)] pub struct ResolvedSearchResult { diff --git a/src/bench/claude.rs b/src/bench/claude.rs index 725360e..d95e005 100644 --- a/src/bench/claude.rs +++ b/src/bench/claude.rs @@ -120,6 +120,16 @@ struct SearchQueryOutput { fts_query: Option, } +#[derive(Deserialize)] +struct FollowupOutput { + #[serde(default)] + sufficient: bool, + #[serde(default)] + queries: Option>, + #[serde(default)] + entities: Vec, +} + #[derive(Deserialize)] struct StoreMemoryCall { summary: String, @@ -375,7 +385,233 @@ impl LlmClient { category: &str, graph_context: &GraphContext, ) -> Result { - // Sort memories chronologically (earliest first) for temporal reasoning + let base = "You are a memory assistant answering questions about long personal \ + conversations between two people who are friends. You receive retrieved memory \ + excerpts from these conversations, each with a relevance score (higher = \ + stronger match) and a date.\n\n\ + Memory types:\n\ + - Numbered memories [1], [2], etc. are directly retrieved for this question.\n\ + - Memories labeled [R1], [R2], etc. are related graph neighbors — contextually \ + connected but not directly matched to the question. Use them as supporting \ + context.\n\n\ + Each memory may include structured metadata:\n\ + - Entities: named entities (people, places, etc.) mentioned in the memory.\n\ + - Topics: subject areas the memory relates to.\n\ + - Mood: emotional tone recorded at the time.\n\ + Use this metadata to verify person references and understand context.\n\n\ + When a memory contains a direct quote or specific detail, prefer that detail \ + over your general knowledge."; + + let instructions = "\n\n\n\ + Base your answer ONLY on the provided memories.\n\n\ + Step-by-step approach:\n\ + 1. Identify which memories contain relevant evidence.\n\ + 2. If the question requires combining information from multiple memories, \ + chain the pieces together.\n\ + 3. Memories are sorted chronologically (earliest first). Time deltas between \ + consecutive memories are pre-computed and shown as '\u{23f1} N days after [X]'. \ + For duration questions, use these pre-computed deltas directly — they are exact. \ + The day of week is shown in each memory's date header.\n\ + 4. For counting questions, enumerate each distinct instance explicitly \ + before stating the final count.\n\n\ + Qualifier verification: if the question names a SPECIFIC role, event \ + type, or entity (e.g., \"marketing director\", \"summer internship\", \ + \"Honda Civic\"), verify that this EXACT qualifier appears in at least \ + one memory. If only a DIFFERENT qualifier exists (e.g., \"sales manager\" \ + instead of \"marketing director\", \"part-time job\" instead of \"summer \ + internship\"), treat the question as unanswerable — do NOT answer from \ + the wrong entity.\n\ + Conversely, if a memory describes the same real-world fact using different \ + words (e.g., \"purchased a standing desk\" answers \"what desk did I \ + buy?\", or \"stayed up until midnight on Tuesday\" answers \"what time \ + did I go to bed the night before Wednesday?\"), extract the answer — do \ + NOT say \"I don't know.\"\n\n\ + Semantic matching: memories may use different words for the same concept \ + (e.g., \"plays guitar\" relates to \"rock music\"). Use the information \ + available even when the phrasing differs.\n\n\ + For hypothetical questions (\"would they...\", \"is it likely...\"), find \ + evidence of relevant preferences or behaviors and commit to a clear answer.\n\n\ + For questions requiring inference or world knowledge (e.g., \"what type \ + of archery bow would suit a beginner?\"), combine memory evidence with \ + your general knowledge to provide the most specific and helpful answer.\n\n\ + For list-type questions (\"what are X's hobbies?\", \"name the places Y \ + visited\"), systematically scan ALL retrieved memories for every relevant \ + item. Collect every instance before answering.\n\n\ + IMPORTANT — do NOT hedge or refuse when evidence exists:\n\ + - Never say \"the memories do not mention\" or \"there is no explicit evidence\" \ + when there IS relevant information in the retrieved memories — even if it \ + requires a small inference.\n\ + - If a memory contains a relevant fact but uses different wording than the \ + question, that IS sufficient evidence. Use it.\n\ + - If multiple memories together imply an answer, synthesize them into a \ + clear response rather than saying the answer is unclear.\n\ + - Prefer giving a well-reasoned answer from available evidence over hedging.\n\ + - Say \"I don't know\" ONLY when truly no relevant evidence exists in ANY \ + of the retrieved memories or their neighbors.\n\n\ + Keep your final answer concise: one or two sentences after your reasoning.\n\ + "; + + self.generate_answer_with_system( + question, + memories, + category, + graph_context, + &format!("{base}{instructions}"), + None, + ) + .await + } + + pub async fn generate_answer_longmemeval( + &self, + question: &str, + memories: &[MemoryContext], + category: &str, + graph_context: &GraphContext, + question_date: Option<&str>, + ) -> Result { + let base = "You are a memory assistant answering questions about a user based on \ + their past chat history with an AI assistant. You receive retrieved memory \ + excerpts from these conversations, each with a relevance score (higher = \ + stronger match) and a date.\n\n\ + Memory types:\n\ + - Numbered memories [1], [2], etc. are directly retrieved for this question.\n\ + - Memories labeled [R1], [R2], etc. are related graph neighbors — contextually \ + connected but not directly matched to the question. Use them as supporting \ + context.\n\n\ + Each memory may include structured metadata:\n\ + - Entities: named entities (people, places, etc.) mentioned in the memory.\n\ + - Topics: subject areas the memory relates to.\n\ + - Mood: emotional tone recorded at the time.\n\ + Use this metadata to verify references and understand context.\n\n\ + When a memory contains a direct quote or specific detail, prefer that detail \ + over your general knowledge."; + + let instructions = "\n\n\n\ + Base your answer ONLY on the provided memories.\n\n\ + Step-by-step approach:\n\ + 1. Identify which memories contain relevant evidence.\n\ + 2. If the question requires combining information from multiple memories, \ + chain the pieces together.\n\ + 3. Memories are sorted chronologically (earliest first). The day of week \ + is shown in each memory's date header. Time deltas between consecutive \ + memories are shown as '\u{23f1} N days after [X]' for reference.\n\ + 4. For relative time references ('two months ago', 'last week', 'recently'), \ + use Today's date (shown above the memories) to calculate which memory the \ + question refers to. Write out the dates and compute the difference step by \ + step. When a memory says 'for N months' or 'since N months ago', compute \ + the implied start date by subtracting from the memory's date.\n\ + 5. For time-related questions, ALWAYS compute dates from the memory headers \ + directly. Extract the relevant dates, compute the difference, and state \ + the result.\n\ + 6. For counting questions ('how many times', 'how many X did I'), follow \ + this EXACT procedure — do NOT skip any step:\n\ + a. Scan ALL memories [1] through [N] AND all neighbors [R1] through [RN]. \ + For each one, ask: does this describe a qualifying instance?\n\ + b. Create a numbered list of every qualifying instance with: memory label, \ + date, and key identifying detail.\n\ + c. After completing the list, do a SECOND pass through ALL memories looking \ + specifically for items you missed — check memories with low relevance \ + scores and graph neighbors, they often contain additional instances.\n\ + d. Check for double-counting: are any two entries the same event described \ + in different memories? Deduplicate.\n\ + e. Only then count the final list.\n\ + Common mistake: stopping after finding 2-3 obvious matches when more exist.\n\n\ + When multiple memories describe the same fact at different times \ + (e.g., the user's job, address, personal record), the MOST RECENT \ + memory is the current truth. Earlier memories are history. Always answer \ + with the latest known value.\n\n\ + Qualifier verification: if the question names a SPECIFIC role, event \ + type, or entity (e.g., \"marketing director\", \"summer internship\", \ + \"Honda Civic\"), verify that this EXACT qualifier appears in at least \ + one memory. If only a DIFFERENT qualifier exists (e.g., \"sales manager\" \ + instead of \"marketing director\", \"part-time job\" instead of \"summer \ + internship\"), treat the question as unanswerable — do NOT answer from \ + the wrong entity.\n\ + Conversely, if a memory describes the same real-world fact using different \ + words (e.g., \"purchased a standing desk\" answers \"what desk did I \ + buy?\", or \"stayed up until midnight on Tuesday\" answers \"what time \ + did I go to bed the night before Wednesday?\"), extract the answer — do \ + NOT say \"I don't know.\"\n\n\ + Semantic matching: memories may use different words for the same concept \ + (e.g., \"plays guitar\" relates to \"rock music\"). Use the information \ + available even when the phrasing differs.\n\n\ + For questions asking for suggestions or recommendations, personalize based on \ + the user's stated preferences, past purchases, habits, and dislikes. Pay \ + attention to negative preferences (things the user wants to avoid). Even \ + memories with low relevance scores may contain user context — read all of \ + them before answering.\n\n\ + For list-type questions (\"how many X have I...\", \"what are my hobbies?\"), \ + systematically scan ALL retrieved memories for every relevant instance. \ + Collect every item before answering.\n\n\ + IMPORTANT — do NOT hedge or refuse when evidence exists:\n\ + - Never say \"the memories do not mention\", \"there is no explicit evidence\", \ + or \"the user did not\" when there IS relevant information — even if the \ + wording differs from the question.\n\ + - The question's phrasing may not match the memory exactly. Different verbs, \ + tenses, or framing do not invalidate a memory. If a memory contains the \ + factual answer, GIVE that answer regardless of wording differences.\n\ + - If information can be INFERRED by combining multiple memories (e.g., \ + computing dates, chaining facts across sessions, adding amounts), provide \ + the inferred answer. An inference from available evidence is NOT the same \ + as making up information.\n\ + - If multiple memories together imply an answer, synthesize them into a \ + clear response rather than saying the answer is unclear.\n\ + - Prefer giving a well-reasoned answer from available evidence over hedging.\n\ + - Say \"I don't know\" or \"I don't have that information\" ONLY when truly \ + no relevant evidence exists in ANY of the retrieved memories or their \ + neighbors. This is the correct response for questions about things the \ + user never discussed.\n\ + - A question is unanswerable ONLY if the specific entity, event, or \ + attribute asked about is genuinely absent from ALL memories. If a memory \ + mentions a specific qualifier, verify that it matches the question before \ + answering or abstaining.\n\n\ + Before concluding \"I don't have that information\", perform these checks:\n\ + 1. ADDITION CHECK: Can two or more memories be added together? (e.g., \ + \"2 items from store A\" + \"3 items from store B\" = 5 total items)\n\ + 2. SYNONYM CHECK: Does a memory describe the same thing using a different \ + verb or noun? (e.g., \"owns X\" answers \"what did I buy?\"; \"went to \ + bed at 2 AM\" answers \"what time did I go to sleep?\")\n\ + 3. CHAIN CHECK: Can the answer be derived by chaining two facts? (e.g., \ + \"event was on Wednesday\" + \"appointment was Thursday\" answers \ + \"what happened the day before my appointment?\")\n\ + If any check succeeds, GIVE the answer — do not refuse.\n\n\ + Answer completeness rules:\n\ + - Include ALL specific details from the source memory: full qualifiers, \ + locations, examples, subcategories. Never summarize away detail.\n\ + - Give the most DEFINITIVE answer the evidence supports. Do NOT hedge with \ + \"close to\", \"approximately\", or \"around\" when a memory states a \ + specific value — state the value directly.\n\ + - For numerical answers, give ONE specific number, not a range.\n\ + - When two items seem tied, pick the one with stronger or more recent \ + evidence rather than reporting a tie.\n\n\ + Output format: State your final answer FIRST in a single line starting \ + with \"ANSWER:\", then explain your reasoning below. This format is \ + mandatory. Example:\n\ + ANSWER: The user's favorite color is blue.\n\ + Reasoning: Memory [3] from 2023-05-14 states...\n\ + "; + + self.generate_answer_with_system( + question, + memories, + category, + graph_context, + &format!("{base}{instructions}"), + question_date, + ) + .await + } + + async fn generate_answer_with_system( + &self, + question: &str, + memories: &[MemoryContext], + _category: &str, + graph_context: &GraphContext, + system: &str, + question_date: Option<&str>, + ) -> Result { let mut sorted_memories = memories.to_vec(); sorted_memories.sort_by_key(|m| m.created_at); @@ -404,7 +640,6 @@ impl LlmClient { label, mem.score, day_of_week, date, mem.text ); - // Precompute time delta from previous memory if i > 0 { let prev = &sorted_memories[i - 1]; let delta_ms = mem.created_at - prev.created_at; @@ -482,75 +717,21 @@ impl LlmClient { ) }; - let base = "You are a memory assistant answering questions about long personal \ - conversations between two people who are friends. You receive retrieved memory \ - excerpts from these conversations, each with a relevance score (higher = \ - stronger match) and a date.\n\n\ - Memory types:\n\ - - Numbered memories [1], [2], etc. are directly retrieved for this question.\n\ - - Memories labeled [R1], [R2], etc. are related graph neighbors — contextually \ - connected but not directly matched to the question. Use them as supporting \ - context.\n\n\ - Each memory may include structured metadata:\n\ - - Entities: named entities (people, places, etc.) mentioned in the memory.\n\ - - Topics: subject areas the memory relates to.\n\ - - Mood: emotional tone recorded at the time.\n\ - Use this metadata to verify person references and understand context.\n\n\ - When a memory contains a direct quote or specific detail, prefer that detail \ - over your general knowledge."; - - let instructions = "\n\n\n\ - Base your answer ONLY on the provided memories.\n\n\ - Step-by-step approach:\n\ - 1. Identify which memories contain relevant evidence.\n\ - 2. If the question requires combining information from multiple memories, \ - chain the pieces together.\n\ - 3. Memories are sorted chronologically (earliest first). Time deltas between \ - consecutive memories are pre-computed and shown as '\u{23f1} N days after [X]'. \ - For duration questions, use these pre-computed deltas directly — they are exact. \ - The day of week is shown in each memory's date header.\n\ - 4. For counting questions, enumerate each distinct instance explicitly \ - before stating the final count.\n\n\ - Person verification: if the question names a specific person, confirm \ - the memory is about THAT person before using it. A memory about one \ - person does not answer a question about a different person.\n\n\ - Semantic matching: memories may use different words for the same concept \ - (e.g., \"plays guitar\" relates to \"rock music\"). Use the information \ - available even when the phrasing differs.\n\n\ - For hypothetical questions (\"would they...\", \"is it likely...\"), find \ - evidence of relevant preferences or behaviors and commit to a clear answer.\n\n\ - For questions requiring inference or world knowledge (e.g., \"what type \ - of archery bow would suit a beginner?\"), combine memory evidence with \ - your general knowledge to provide the most specific and helpful answer.\n\n\ - For list-type questions (\"what are X's hobbies?\", \"name the places Y \ - visited\"), systematically scan ALL retrieved memories for every relevant \ - item. Collect every instance before answering.\n\n\ - IMPORTANT — do NOT hedge or refuse when evidence exists:\n\ - - Never say \"the memories do not mention\" or \"there is no explicit evidence\" \ - when there IS relevant information in the retrieved memories — even if it \ - requires a small inference.\n\ - - If a memory contains a relevant fact but uses different wording than the \ - question, that IS sufficient evidence. Use it.\n\ - - If multiple memories together imply an answer, synthesize them into a \ - clear response rather than saying the answer is unclear.\n\ - - Prefer giving a well-reasoned answer from available evidence over hedging.\n\ - - Say \"I don't know\" ONLY when truly no relevant evidence exists in ANY \ - of the retrieved memories or their neighbors.\n\n\ - Keep your final answer concise: one or two sentences after your reasoning.\n\ - "; - - let system = format!("{base}{instructions}"); + let reasoning_hint = "Think step by step. First identify which memories are relevant. Then extract the key facts. If the question involves dates, write out the dates and compute. If it involves counting, list each instance. Finally, state your answer."; - let reasoning_hint = "Identify the relevant evidence from the memories, then answer."; + let date_context = match question_date { + Some(d) if !d.is_empty() => format!("Today's date: {d}\n\n"), + _ => String::new(), + }; let user = format!( - "Retrieved memories:\n{context}{neighbor_context}\n\n\ + "{date_context}Retrieved memories:\n{context}{neighbor_context}\n\n\ Question: {question}\n\n\ {reasoning_hint}\n\n\ Answer:" ); - self.call(&system, &user, 4096).await + self.call(system, &user, 4096).await } pub async fn judge_answer( @@ -638,6 +819,113 @@ impl LlmClient { } } + pub async fn judge_answer_longmemeval( + &self, + question: &str, + gold_answer: &str, + generated_answer: &str, + question_type: &str, + is_abstention: bool, + ) -> Result { + let prompt = if is_abstention { + format!( + "I will give you an unanswerable question, an explanation, and a response from \ + a model. Please answer yes if the model correctly identifies the question as \ + unanswerable. The model could say that the information is incomplete, or some \ + other information is given but the asked information is not.\n\n\ + Question: {question}\n\n\ + Explanation: {gold_answer}\n\n\ + Model Response: {generated_answer}\n\n\ + Does the model correctly identify the question as unanswerable? Answer yes or \ + no only." + ) + } else { + match question_type { + "single-session-user" | "single-session-assistant" | "multi-session" => { + format!( + "I will give you a question, a correct answer, and a response from a \ + model. Please answer yes if the response contains the correct answer. \ + Otherwise, answer no. If the response is equivalent to the correct \ + answer or contains all the intermediate steps to get the correct answer, \ + you should also answer yes. If the response only contains a subset of \ + the information required by the answer, answer no. \n\n\ + Question: {question}\n\n\ + Correct Answer: {gold_answer}\n\n\ + Model Response: {generated_answer}\n\n\ + Is the model response correct? Answer yes or no only." + ) + } + "temporal-reasoning" => { + format!( + "I will give you a question, a correct answer, and a response from a \ + model. Please answer yes if the response contains the correct answer. \ + Otherwise, answer no. If the response is equivalent to the correct \ + answer or contains all the intermediate steps to get the correct answer, \ + you should also answer yes. If the response only contains a subset of \ + the information required by the answer, answer no. In addition, do not \ + penalize off-by-one errors for the number of days. If the question asks \ + for the number of days/weeks/months, etc., and the model makes off-by-one \ + errors (e.g., predicting 19 days when the answer is 18), the model's \ + response is still correct. \n\n\ + Question: {question}\n\n\ + Correct Answer: {gold_answer}\n\n\ + Model Response: {generated_answer}\n\n\ + Is the model response correct? Answer yes or no only." + ) + } + "knowledge-update" => { + format!( + "I will give you a question, a correct answer, and a response from a \ + model. Please answer yes if the response contains the correct answer. \ + Otherwise, answer no. If the response contains some previous information \ + along with an updated answer, the response should be considered as \ + correct as long as the updated answer is the required answer.\n\n\ + Question: {question}\n\n\ + Correct Answer: {gold_answer}\n\n\ + Model Response: {generated_answer}\n\n\ + Is the model response correct? Answer yes or no only." + ) + } + "single-session-preference" => { + format!( + "I will give you a question, a rubric for desired personalized response, \ + and a response from a model. Please answer yes if the response satisfies \ + the desired response. Otherwise, answer no. The model does not need to \ + reflect all the points in the rubric. The response is correct as long as \ + it recalls and utilizes the user's personal information correctly.\n\n\ + Question: {question}\n\n\ + Rubric: {gold_answer}\n\n\ + Model Response: {generated_answer}\n\n\ + Is the model response correct? Answer yes or no only." + ) + } + _ => { + format!( + "I will give you a question, a correct answer, and a response from a \ + model. Please answer yes if the response contains the correct answer. \ + Otherwise, answer no. If the response is equivalent to the correct \ + answer or contains all the intermediate steps to get the correct answer, \ + you should also answer yes. If the response only contains a subset of \ + the information required by the answer, answer no. \n\n\ + Question: {question}\n\n\ + Correct Answer: {gold_answer}\n\n\ + Model Response: {generated_answer}\n\n\ + Is the model response correct? Answer yes or no only." + ) + } + } + }; + + let response = self + .call("You are an evaluation assistant.", &prompt, 10) + .await?; + let label = response.to_lowercase().contains("yes"); + Ok(JudgeResult { + correct: label, + reason: response.trim().to_string(), + }) + } + pub async fn process_turn( &self, recent_turns: &[String], @@ -645,7 +933,36 @@ impl LlmClient { recent_memories: &[(String, String)], turn_timestamp_ms: i64, ) -> Result, String> { - let system = "You are an AI assistant listening to a conversation between two people. \ + Self::process_turn_inner( + self, + recent_turns, + current_turn, + recent_memories, + turn_timestamp_ms, + Self::SYSTEM_PROMPT_LOCOMO, + ) + .await + } + + pub async fn process_turn_longmemeval( + &self, + recent_turns: &[String], + current_turn: &str, + recent_memories: &[(String, String)], + turn_timestamp_ms: i64, + ) -> Result, String> { + Self::process_turn_inner( + self, + recent_turns, + current_turn, + recent_memories, + turn_timestamp_ms, + Self::SYSTEM_PROMPT_LONGMEMEVAL, + ) + .await + } + + const SYSTEM_PROMPT_LOCOMO: &str = "You are an AI assistant listening to a conversation between two people. \ You have a memory tool called store_memory that saves information for later recall.\n\n\ You will see recent conversation context followed by the LATEST turn. \ Store new factual information revealed in the latest turn, but be highly selective — \ @@ -747,6 +1064,133 @@ impl LlmClient { Aim for precision: fewer high-quality memories are better than many redundant ones.\n\ "; + const SYSTEM_PROMPT_LONGMEMEVAL: &str = "You are an AI assistant reviewing a chat history between a user and an \ + AI assistant. You have a memory tool called store_memory that saves information \ + about the user for later recall.\n\n\ + You will see recent conversation context followed by the LATEST turn. \ + Store new factual information about the user revealed in the latest turn, but \ + be highly selective — only store facts that would be useful to answer future \ + questions about the user.\n\n\ + \n\ + Your job: extract genuinely new, specific, useful facts from the latest turn.\n\ + Store facts from BOTH the user AND the assistant:\n\ + - User turns: what the user says about themselves, their life, preferences, actions.\n\ + - Assistant turns: specific facts, recommendations, data, or answers the assistant \ + provided. The user may later ask \"what did you tell me about X?\" or \"remind me \ + of that thing you mentioned.\" These must be retrievable.\n\ + Quality over quantity. A typical USER turn contains 0-2 memories worth storing. \ + ASSISTANT turns with rich content (lists, schedules, detailed recommendations, \ + creative works) may warrant 2-4 memories to capture distinct details. \ + Most turns with casual chat, reactions, or continuation of an already-stored topic \ + should produce zero new memories.\n\ + \n\n\ + \n\ + Store facts that might be asked about later:\n\ + FROM THE USER:\n\ + - Concrete biographical details: name, relationships, occupation, where they live\n\ + - Specific preferences: favorite foods, hobbies, media they enjoy, tools they use\n\ + - Events and plans: trips, milestones, upcoming events with dates\n\ + - Quantities and durations: \"5 years\", \"every morning\", \"three times a week\", \"$200\"\n\ + - Exact counts: how many children, pets, siblings, trips, items purchased, etc.\n\ + - Purchases and acquisitions: what the user bought, where, how much, specific models\n\ + - People the user mentions: friends, family, coworkers — store their names and \ + relationship to the user.\n\ + - Specific names, titles, dates, objects: pets' names, book titles, brand names\n\ + - Emotional reactions and feelings about events or experiences\n\ + - Activities and hobbies: what the user does, how often, with whom\n\ + - Completion events: when the user finishes a book, course, project, or reaches a \ + goal, store the completion with the date. These are critical for duration questions.\n\ + - Changes in the user's situation: new job, moved, changed preferences. When a \ + fact changes, the new memory MUST supersede the old one.\n\n\ + FROM THE ASSISTANT:\n\ + - Specific recommendations: app names, product names, restaurant names, tools \ + suggested to the user (e.g., \"the assistant recommended Memrise for language learning\").\n\ + - Factual data provided: statistics, study results, specific numbers, dates, or \ + measurements the assistant cited (e.g., \"the assistant cited a study with 38 subjects\").\n\ + - Explanations of specific concepts the user asked about.\n\ + - Schedules, plans, or structured information the assistant created for the user \ + (e.g., shift rotations, itineraries, meal plans). Store individual assignments, \ + not just the fact that a schedule was created.\n\ + - Exact quotes, literary references, definitions, or terminology the assistant \ + provided. If the assistant quoted a source, listed alternative terms, or cited a \ + specific passage, store the verbatim content in full_text.\n\ + - Creative content the assistant produced: song lyrics with chords, poems, stories. \ + Store enough detail that the user could ask about specific parts later.\n\ + - Named individuals mentioned in factual discussions (e.g., a scientist, advisor, \ + or author the assistant referenced by name).\n\ + - Any specific detail the user might later say \"remind me what you said about...\" for.\n\ + \n\n\ + \n\ + Use EXACT nouns from the conversation. Store the precise name, title, or object: \ + \"To Kill a Mockingbird\" not \"a book\", \"sneakers\" not \"shoes\", \"Portugal\" not \ + \"a European country\", \"labrador\" not \"dog\".\n\ + Always include specific items, places, brands, species, colors, and quantities.\n\ + If a book, movie, song, or other titled work is discussed, include the exact title \ + in both summary and full_text.\n\ + Resolve relative times ('last week', 'yesterday') to approximate absolute dates \ + using the conversation timestamp. The resolved date MUST appear in the summary \ + text itself (e.g., 'On approximately May 3, 2023, the user mentioned...'). Do NOT \ + put dates only in metadata — the summary must contain the date so it is searchable \ + and self-contained when retrieved independently.\n\ + \n\n\ + \n\ + Every summary should refer to the user as \"the user\" consistently. Each memory \ + is retrieved independently without surrounding context, so it must be completely \ + self-contained and unambiguous.\n\ + BAD: \"They mentioned taking calligraphy lessons.\"\n\ + GOOD: \"The user mentioned taking calligraphy lessons since approximately March 2023.\"\n\ + BAD: \"Got a new job at the hospital.\"\n\ + GOOD: \"The user started a new job at Memorial Hospital in May 2023.\"\n\ + When the user mentions other people by name, include both \"the user\" and the \ + named person as entities: \"The user said that Rachel moved to the suburbs.\"\n\ + \n\n\ + \n\ + CRITICAL: Before creating any memory, check the already-stored memories listed below.\n\ + If a fact is already captured — even partially, with different wording, or as part of \ + a broader memory — skip it. Only store if the turn reveals a genuinely NEW fact \ + that no existing memory covers.\n\ + If an existing memory is OUTDATED or INCOMPLETE and the conversation reveals new \ + information that changes it, store a new memory with \"supersedes\" set to the \ + ID of the old memory. This is especially important for facts that change over time \ + (e.g., the user's current job, address, personal best time, number of items owned).\n\ + \n\n\ + \n\ + The tags you attach (entities, topics, emotions) are what the search system \ + uses to retrieve memories later. Thorough, accurate tagging directly determines \ + whether a memory can be found. If a fact involves a person, tag them as an entity. \ + If it has emotional content, tag the emotion. If it relates to a topic, tag it.\n\n\ + store_memory accepts:\n\ + - \"summary\" (required, string): A concise 1-2 sentence summary of the key fact.\n\ + - \"full_text\" (optional, string): Detailed version with all context, quotes, and \ + specifics. Provide this for any memory containing specific names, titles, dates, \ + numbers, or quotes. Include exact names, dates, numbers, and nuance.\n\ + - \"entities\" (required, array of strings): ALL people, pets, places, organizations, \ + titles, and proper nouns in this memory. Always include \"user\" as an entity. \ + Use canonical names (e.g., \"Rachel\", \"Portugal\", \"Spotify\").\n\ + - \"topics\" (required, array of strings): 1-5 topic keywords. Examples: \"adoption\", \ + \"career\", \"cooking\", \"music\", \"travel\". Use lowercase.\n\ + - \"emotions\" (optional, array of strings): Only include if clear emotional content \ + is present. Examples: \"happy\", \"anxious\", \"grateful\".\n\ + - \"supersedes\" (optional, string): ID of an existing memory this one replaces.\n\n\ + Output: {\"store_memory\": [...]} or {\"store_memory\": []} if nothing new to store.\n\ + Output ONLY valid JSON. No markdown fences, no explanation, no commentary.\n\ + \n\n\ + \n\ + Skip greetings, filler, emotional reactions without new factual content, and \ + restatements of already-stored information.\n\ + Each memory must capture a UNIQUE fact not present in any existing memory.\n\ + If the latest turn just continues discussing a topic already stored, store nothing.\n\ + Aim for precision: fewer high-quality memories are better than many redundant ones.\n\ + "; + + async fn process_turn_inner( + &self, + recent_turns: &[String], + current_turn: &str, + recent_memories: &[(String, String)], + turn_timestamp_ms: i64, + system: &str, + ) -> Result, String> { let mut user_msg = String::new(); user_msg.push_str(&format!( "Conversation timestamp: {}\n\n", @@ -817,22 +1261,28 @@ impl LlmClient { one subject area. At most ONE topic (multiple use AND logic, which is too restrictive).\n\ - \"emotions\" (optional, array of strings): Filter by emotional tag. Only include when \ the question asks about feelings. At most ONE emotion.\n\ - - \"depth\" (optional, integer 0-3, default 2): Graph hops for related memories. \ - Use 2 for opinions, hypotheticals, inference questions. Use 1 for simple factual \ - lookups. Use 3 for complex multi-hop reasoning.\n\ + - \"depth\" (optional, integer 0-2, default 2): Graph hops for related memories. \ + Use 0 for simple factual lookups with a clear entity. Use 1 for straightforward \ + single-fact questions. Use 2 (default) for temporal, multi-step, or inference \ + questions.\n\ - \"time_range_start\" (optional, integer): Lower bound timestamp in milliseconds \ since Unix epoch.\n\ - \"time_range_end\" (optional, integer): Upper bound timestamp in milliseconds \ since Unix epoch.\n\ \n\n\ \n\ - Use multiple queries when the question has multiple angles:\n\ - - \"What movies has Alice watched?\" -> one query about Alice's movie-watching habits, \ - another about specific film titles Alice mentioned.\n\ - - \"Would Marcus enjoy cooking?\" -> one query about Marcus's food-related activities, \ - another about Marcus's hobbies or kitchen interests.\n\ - For simple factual questions, one query is sufficient.\n\ + ALWAYS use 2-3 queries. Multiple queries from different angles dramatically \ + improve recall. Each query should approach the question differently:\n\ + - \"What movies has Alice watched?\" -> (1) Alice's movie-watching habits, \ + (2) specific film titles Alice mentioned, (3) Alice entertainment preferences.\n\ + - \"Would Marcus enjoy cooking?\" -> (1) Marcus's food-related activities, \ + (2) Marcus's hobbies or interests, (3) things Marcus has said about cooking.\n\ + - \"What's my cat's name?\" -> (1) user's pet cat name, (2) user's animals \ + or pets at home.\n\ Each query should approach from a genuinely different angle, not be a paraphrase.\n\ + For counting or aggregation questions ('how many X', 'how much total'), use \ + 3 queries to cover different instances — each instance may be stored in \ + a separate memory with different wording.\n\ For opinion or inference questions, search for underlying facts and experiences \ rather than the inference itself.\n\ \n\n\ @@ -872,7 +1322,7 @@ impl LlmClient { entities: parsed.entities, topics: parsed.topics, emotions: parsed.emotions, - depth: parsed.depth.unwrap_or(2).min(3), + depth: parsed.depth.unwrap_or(2).min(2), time_range_start: parsed.time_range_start, time_range_end: parsed.time_range_end, }) @@ -892,6 +1342,80 @@ impl LlmClient { } } + pub async fn construct_followup_query( + &self, + question: &str, + retrieved_summaries: &str, + ) -> Result, String> { + let system = "You are evaluating whether retrieved memories are sufficient to answer a \ + question. You receive the question and summaries of what was already retrieved.\n\n\ + Your job is to find GAPS and chase CONNECTIONS across sessions.\n\n\ + Step 1: Read the question carefully. What specific facts are needed to answer it?\n\ + Step 2: Check the retrieved memories. Which needed facts are present? Which are \ + missing?\n\ + Step 3: Look at entities, people, places, and topics MENTIONED in the retrieved \ + memories. Are there related memories that could be found by searching for those \ + entities in a different context?\n\n\ + Common patterns requiring follow-up:\n\ + - Counting/aggregation: some items found, but the question implies more exist. \ + Search for related items using different wording or categories.\n\ + - Cross-session connections: a retrieved memory mentions a person, place, or \ + event that appears in other sessions. Search for that entity specifically.\n\ + - Temporal chains: the question asks about changes over time but only one \ + time point was found. Search for earlier or later mentions.\n\ + - The question asks about multiple aspects but only one was found.\n\ + - All retrieved memories have very low relevance (none clearly match).\n\n\ + If the retrieved memories are sufficient, output:\n\ + {\"sufficient\": true}\n\n\ + If more information is needed, output a JSON object with:\n\ + - \"sufficient\": false\n\ + - \"queries\" (array of 1-3 objects with \"query\" and optional \"fts_query\"): \ + Search queries targeting SPECIFICALLY what is missing. Use entities and topics \ + from the retrieved memories to find connected information in other sessions. \ + Do NOT repeat the original queries — search for the gap.\n\ + - \"entities\" (optional, array): Entity filters for the follow-up.\n\n\ + Default to searching MORE rather than less — a follow-up that finds nothing \ + costs little, but a skipped follow-up that would have found the answer is \ + a missed opportunity.\n\n\ + Output ONLY the JSON object. No markdown fences, no explanation."; + + let user_msg = format!("Question: {question}\n\nAlready retrieved:\n{retrieved_summaries}"); + let response = self.call(system, &user_msg, 512).await?; + let cleaned = response + .trim() + .trim_start_matches("```json") + .trim_start_matches("```") + .trim_end_matches("```") + .trim(); + + match serde_json::from_str::(cleaned) { + Ok(parsed) => { + if parsed.sufficient || parsed.queries.is_none() { + return Ok(None); + } + let queries = parsed + .queries + .unwrap() + .into_iter() + .map(|q| SearchQuery { + query: q.query, + fts_query: q.fts_query, + }) + .collect(); + Ok(Some(SearchParams { + queries, + entities: parsed.entities, + topics: Vec::new(), + emotions: Vec::new(), + depth: 2, + time_range_start: None, + time_range_end: None, + })) + } + Err(_) => Ok(None), + } + } + async fn call(&self, system: &str, user: &str, max_tokens: u32) -> Result { let messages = vec![ChatMessage { role: "user".to_string(), @@ -953,13 +1477,11 @@ impl LlmClient { }], }, generation_config: GeminiGenerationConfig { - max_output_tokens: max_tokens, + max_output_tokens: max_tokens + 2048, temperature: 0.0, - thinking_config: if self.model.contains("flash") { - Some(GeminiThinkingConfig { thinking_budget: 0 }) - } else { - None - }, + thinking_config: Some(GeminiThinkingConfig { + thinking_budget: 2048, + }), }, }; self.client diff --git a/src/bench/locomo.rs b/src/bench/locomo.rs index 607f1f7..26bdf6c 100644 --- a/src/bench/locomo.rs +++ b/src/bench/locomo.rs @@ -306,13 +306,15 @@ fn parse_dataset(path: &Path) -> Result, Box, } impl BenchHarness { - async fn new(base_config: &RecalldConfig) -> Result> { + pub(super) async fn new( + base_config: &RecalldConfig, + ) -> Result> { let temp_dir = tempfile::TempDir::new()?; let mut config = base_config.clone(); @@ -328,9 +330,34 @@ impl BenchHarness { Ok(Self { system, - _temp_dir: temp_dir, + _temp_dir: Some(temp_dir), }) } + + pub(super) async fn new_persistent( + base_config: &RecalldConfig, + dir: &std::path::Path, + ) -> Result> { + std::fs::create_dir_all(dir)?; + let mut config = base_config.clone(); + config.storage.data_dir = dir + .to_str() + .ok_or("cache dir path is not valid UTF-8")? + .to_string(); + config.decay.sweep_interval_hours = 999_999.0; + config.decay.disable_sweep = true; + let system = Recalld::new(config).await?; + Ok(Self { + system, + _temp_dir: None, + }) + } + + pub(super) fn has_memories(&self) -> bool { + let storage = self.system.storage(); + let guard = storage.read().expect("storage lock poisoned"); + guard.count().unwrap_or(0) > 0 + } } // ── Results ─────────────────────────────────────────────────────── @@ -1203,7 +1230,7 @@ async fn run_conversation( // ── Ingest ──────────────────────────────────────────────────────── -async fn store_single_memory( +pub(super) async fn store_single_memory( harness: &BenchHarness, summary: &str, full_text: Option<&str>, @@ -1465,16 +1492,16 @@ async fn ingest_conversation( // ── Search ──────────────────────────────────────────────────────── #[derive(Clone)] -struct ScoredMemory { - memory_id: MemoryId, - text: String, - full_text: Option, - score: f32, - created_at: i64, - tags: Vec, +pub(super) struct ScoredMemory { + pub(super) memory_id: MemoryId, + pub(super) text: String, + pub(super) full_text: Option, + pub(super) score: f32, + pub(super) created_at: i64, + pub(super) tags: Vec, } -fn load_full_texts( +pub(super) fn load_full_texts( harness: &BenchHarness, results: Vec, ) -> Result, Box> { @@ -1517,7 +1544,7 @@ fn load_full_texts( Ok(scored) } -async fn search_memories( +pub(super) async fn search_memories( harness: &BenchHarness, question: &str, top_k: usize, @@ -1552,7 +1579,7 @@ async fn search_memories( params.queries, params.entities, tags, - params.depth.min(3) as u8, + params.depth.min(2) as u8, params.time_range_start, params.time_range_end, ) @@ -1604,11 +1631,7 @@ async fn search_memories( let _ = writeln!(debug_log, " FTS[{}]: {:?}", qi + 1, fq); } - let per_query_limit = if queries.len() > 1 { - top_k / 2 + 5 - } else { - top_k - }; + let per_query_limit = top_k; let filter = PipelineSearchFilter { require_tags: require_tags.clone(), ..PipelineSearchFilter::default() @@ -1662,12 +1685,115 @@ async fn search_memories( ); } + // Iterative retrieval: up to 2 follow-up rounds. + if let Some(llm) = search_llm { + for round in 1..=1 { + let summaries: String = results + .iter() + .enumerate() + .map(|(i, r)| { + let snippet: String = r.text.chars().take(150).collect(); + format!("[{}] {}", i + 1, snippet) + }) + .collect::>() + .join("\n"); + + match llm.construct_followup_query(question, &summaries).await { + Ok(Some(followup)) => { + let _ = writeln!(debug_log, " -- Follow-up round {round} --"); + let followup_tags: Vec = followup + .entities + .iter() + .filter_map(|e| { + crate::model::Tag::new(format!("entity/{}", e.to_lowercase())).ok() + }) + .collect(); + + let mut followup_merged: HashMap = + results.iter().map(|r| (r.memory_id, r.clone())).collect(); + let prev_count = followup_merged.len(); + + for (qi, bq) in followup.queries.iter().enumerate() { + let _ = + writeln!(debug_log, " FU{round}-Query[{}]: {:?}", qi + 1, bq.query); + if let Some(ref fq) = bq.fts_query { + let _ = writeln!(debug_log, " FU{round}-FTS[{}]: {:?}", qi + 1, fq); + } + let filter = PipelineSearchFilter { + require_tags: followup_tags.clone(), + ..PipelineSearchFilter::default() + }; + let query = SearchQuery { + text: Some(bq.query.clone()), + fts_query: bq.fts_query.clone(), + namespace: "default".to_string(), + filter, + limit: top_k, + min_score: 0.0, + include_ghosts: false, + query_mode: QueryMode::EmbeddingPlusMetadata, + graph_depth: 2, + time_range_start: None, + time_range_end: None, + entities: followup.entities.clone(), + }; + if let Ok(response) = harness.system.query_engine().search(query).await { + if let Ok(new_results) = load_full_texts(harness, response.results) { + for r in new_results { + followup_merged + .entry(r.memory_id) + .and_modify(|existing| { + if r.score > existing.score { + *existing = r.clone(); + } + }) + .or_insert(r); + } + } + } + } + + let new_unique = followup_merged.len() - prev_count; + results = followup_merged.into_values().collect(); + results.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + results.truncate(top_k); + + let _ = writeln!( + debug_log, + " Follow-up round {round}: {new_unique} new, {} total after merge", + results.len() + ); + + if new_unique == 0 { + let _ = writeln!(debug_log, " No new results, stopping follow-ups"); + break; + } + } + Ok(None) => { + let _ = writeln!(debug_log, " Follow-up round {round}: not needed"); + break; + } + Err(e) => { + let _ = writeln!(debug_log, " Follow-up round {round}: error: {e}"); + break; + } + } + } + } + Ok(results) } /// Collect graph context: 1-hop neighbor summaries + edge relationships /// between all memories (results and neighbors). -async fn collect_graph_context(harness: &BenchHarness, results: &[ScoredMemory]) -> GraphContext { +pub(super) async fn collect_graph_context( + harness: &BenchHarness, + results: &[ScoredMemory], +) -> GraphContext { let result_ids: HashSet = results.iter().map(|m| m.memory_id).collect(); // Build label map for results. diff --git a/src/cache/warming_adapters.rs b/src/cache/warming_adapters.rs index e29af6b..ab87683 100644 --- a/src/cache/warming_adapters.rs +++ b/src/cache/warming_adapters.rs @@ -39,7 +39,7 @@ impl StorageEngine for WarmingStorageAdapter { match storage_r.get_record(id) { Ok(Some(disk)) => Ok(Some(CachedRecord::from(&disk))), Ok(None) => Ok(None), - Err(e) => Err(anyhow::anyhow!("{}", e)), + Err(e) => Err(anyhow::Error::from(e)), } } @@ -55,7 +55,7 @@ impl StorageEngine for WarmingStorageAdapter { match storage_r.get_vector(namespace_id, vector_slot) { Ok(Some(v)) => Ok(v), Ok(None) => Err(anyhow::anyhow!("vector not found")), - Err(e) => Err(anyhow::anyhow!("{}", e)), + Err(e) => Err(anyhow::Error::from(e)), } } @@ -69,7 +69,7 @@ impl StorageEngine for WarmingStorageAdapter { .map_err(|e| anyhow::anyhow!("storage lock poisoned: {e}"))?; storage_r .get_outgoing_edges(id) - .map_err(|e| anyhow::anyhow!("{}", e)) + .map_err(anyhow::Error::from) } } @@ -99,6 +99,6 @@ impl EdgeStore for StorageEdgeAdapter { .map_err(|e| anyhow::anyhow!("storage lock poisoned: {e}"))?; storage_r .get_outgoing_edges(id) - .map_err(|e| anyhow::anyhow!("{}", e)) + .map_err(anyhow::Error::from) } } diff --git a/src/decay/fsrs.rs b/src/decay/fsrs.rs index 983621e..62f46b8 100644 --- a/src/decay/fsrs.rs +++ b/src/decay/fsrs.rs @@ -134,18 +134,25 @@ const W6: f32 = 0.75; #[allow(dead_code)] const W7: f32 = 0.6; -/// FSRS v4.5 default parameters w11-w14. +/// FSRS v4.5 default parameters w11-w14 (post-lapse stability formula). /// -/// w11-w13: Stability-after-failure formula (not used -- Recalld -/// has no failure signal; all accesses are treated as successful recall). -/// w14: Scaling factor in FSRS-5 difficulty update (not used in v4.5 mode). -#[allow(dead_code)] +/// Used by `review_stability` when quality = 1 (forgot) to compute +/// post-lapse stability via the FSRS v4.5 formula: +/// +/// ```text +/// S'_f = w11 * D^(-w12) * ((S + 1)^w13 - 1) * e^(w14 * (1 - R)) +/// ``` +/// +/// With default values (w12 = 0.0, w14 = 0.0), this simplifies to: +/// S'_f = w11 * ((S + 1)^w13 - 1) +/// +/// This produces an absolute stability value that replaces (rather than +/// multiplies) the old stability, modeling the empirical finding that +/// failing to recall a memory resets it to a low-but-nonzero stability +/// that depends on how strong the memory was before the lapse. const W11: f32 = 1.2; -#[allow(dead_code)] const W12: f32 = 0.0; -#[allow(dead_code)] const W13: f32 = 0.3261; -#[allow(dead_code)] const W14: f32 = 0.0; /// FSRS v4.5 default parameters w15-w19. @@ -590,24 +597,30 @@ impl<'a> FsrsEngine<'a> { /// update without requiring a full `DecayState`. /// /// Quality ratings map to FSRS behavior: - /// - 1 (forgot): stability decreases (lapse — uses the SInc formula - /// with a penalty factor to reduce stability) + /// - 1 (forgot): stability decreases (lapse — uses the FSRS v4.5 + /// post-lapse formula to compute a new, lower stability) /// - 2 (hard): moderate stability increase /// - 3 (good): standard stability increase (default FSRS recall) /// - 4 (easy): strongest stability increase /// - /// The core FSRS recall formula is: + /// For quality 2-4, the core FSRS recall formula is: /// ```text /// SInc = e^W8 * (11 - D) * S^(-W9) * (e^(W10 * (1 - R)) - 1) + 1 /// new_stability = old_stability * SInc /// ``` /// /// For quality != 3 (Good), a grade multiplier adjusts SInc: - /// - Quality 1 (forgot): multiplier = 0.2 (significant penalty) /// - Quality 2 (hard): multiplier = 0.7 /// - Quality 3 (good): multiplier = 1.0 (standard FSRS) /// - Quality 4 (easy): multiplier = 1.3 /// + /// For quality 1 (forgot), the FSRS v4.5 post-lapse formula is used: + /// ```text + /// S'_f = w11 * D^(-w12) * ((S + 1)^w13 - 1) * e^(w14 * (1 - R)) + /// ``` + /// This produces an absolute stability value (typically much lower + /// than the old stability), modeling the memory reset after a lapse. + /// /// # Arguments /// /// - `stability`: Current FSRS stability S, in days. @@ -628,12 +641,27 @@ impl<'a> FsrsEngine<'a> { // Compute current retrievability at the moment of review. let current_r = self.retrievability(elapsed_days, stability, decay_rate_multiplier); - // Compute the base SInc using the standard FSRS formula. + if quality == 1 { + // Quality 1 (forgot): use the FSRS v4.5 post-lapse formula. + // + // S'_f = w11 * D^(-w12) * ((S + 1)^w13 - 1) * e^(w14 * (1 - R)) + // + // This produces a new absolute stability value (not a multiplier), + // typically much lower than the old stability. For example, with + // S = 3.7145, the formula yields S'_f ~ 0.76. + let d = self.config.difficulty; + let new_stability = W11 + * d.powf(-W12) + * ((stability + 1.0).powf(W13) - 1.0) + * (W14 * (1.0 - current_r)).exp(); + + return new_stability.clamp(STABILITY_FLOOR, STABILITY_CEILING); + } + + // Quality 2-4: use the multiplicative SInc approach. let base_sinc = self.stability_increase(stability, current_r, self.config.difficulty); - // Apply a grade-dependent multiplier to the growth portion (SInc - 1). let grade_multiplier = match quality { - 1 => 0.2, // forgot — substantial penalty 2 => 0.7, // hard — reduced growth 3 => 1.0, // good — standard FSRS _ => 1.3, // easy (4+) — boosted growth diff --git a/src/decay/mod.rs b/src/decay/mod.rs index 5078008..5c812ca 100644 --- a/src/decay/mod.rs +++ b/src/decay/mod.rs @@ -28,6 +28,7 @@ pub use sweep::{ DecaySweepRunner, PhaseTransition, SweepConfig, SweepConfigError, SweepRecordError, SweepResult, }; +use crate::model::constants::ACCESS_HISTORY_MAX; use crate::model::{AccessEvent, DecayPhase, MemoryId}; /// Mutable decay-related fields on a memory record. @@ -76,15 +77,12 @@ pub struct DecayState { } impl DecayState { - /// Maximum number of access events retained in the history ring buffer. - pub const MAX_ACCESS_HISTORY: usize = 32; - /// Push an access event into the bounded history ring buffer. /// - /// If the buffer is full (>= MAX_ACCESS_HISTORY), the oldest entry + /// If the buffer is full (>= ACCESS_HISTORY_MAX), the oldest entry /// is removed before pushing. pub fn push_access(&mut self, event: AccessEvent) { - if self.access_history.len() >= Self::MAX_ACCESS_HISTORY { + if self.access_history.len() >= ACCESS_HISTORY_MAX { self.access_history.remove(0); } self.access_history.push(event); diff --git a/src/decay/storage_adapter.rs b/src/decay/storage_adapter.rs index f95cb1a..5d009f6 100644 --- a/src/decay/storage_adapter.rs +++ b/src/decay/storage_adapter.rs @@ -73,10 +73,18 @@ pub fn update_decay_state( id: MemoryId, phase: DecayPhase, strength: f32, + decay_strength: f32, stability: f32, is_permastore: bool, ) -> Result<(), SweepError> { storage - .update_decay_state(id, phase, strength, stability, is_permastore) + .update_decay_state( + id, + phase, + strength, + decay_strength, + stability, + is_permastore, + ) .map_err(|e| SweepError::Storage(e.to_string())) } diff --git a/src/decay/sweep.rs b/src/decay/sweep.rs index d630c90..6f33844 100644 --- a/src/decay/sweep.rs +++ b/src/decay/sweep.rs @@ -265,6 +265,9 @@ pub enum SweepError { struct PendingTransition { memory_id: MemoryId, from_phase: DecayPhase, + /// Raw FSRS retrievability (without connection bonus). + raw_r: f32, + /// Effective retrievability (with connection bonus). effective_r: f32, /// FSRS stability from the scan-time metadata snapshot. stability: f32, @@ -485,7 +488,7 @@ impl DecaySweepRunner { /// in the same sweep. /// 2. The cheapest operations (phase 3 records are smallest) run first. #[instrument(skip_all)] - async fn execute_sweep( + pub(crate) async fn execute_sweep( config: &SweepConfig, decay_config: &DecayConfig, activation_config: &ActivationConfig, @@ -493,9 +496,33 @@ impl DecaySweepRunner { graph: &SharedGraph, cache: &Arc, global_decay_multiplier: f64, + ) -> SweepResult { + Self::execute_sweep_at( + config, + decay_config, + activation_config, + storage, + graph, + cache, + global_decay_multiplier, + None, + ) + .await + } + + #[instrument(skip_all)] + pub(crate) async fn execute_sweep_at( + config: &SweepConfig, + decay_config: &DecayConfig, + activation_config: &ActivationConfig, + storage: &Arc>, + graph: &SharedGraph, + cache: &Arc, + global_decay_multiplier: f64, + as_of_millis: Option, ) -> SweepResult { let start = Instant::now(); - let now_millis = chrono::Utc::now().timestamp_millis(); + let now_millis = as_of_millis.unwrap_or_else(|| chrono::Utc::now().timestamp_millis()); let engine = FsrsEngine::new(decay_config); let mut result = SweepResult::default(); @@ -689,7 +716,7 @@ impl DecaySweepRunner { }; // Evaluate this record for a possible phase transition. - let effective_r = Self::evaluate_record( + let retrievability = Self::evaluate_record( *memory_id, phase, &meta, @@ -710,11 +737,12 @@ impl DecaySweepRunner { pending_transitions.clear(); } - // Update decay_strength in storage and sync graph state. - // Skip if the record was not evaluated (permastore/decay-disabled). - if let Some(eff_r) = effective_r { + // Update strength and decay_strength in storage and sync + // graph state. Skip if the record was not evaluated + // (permastore/decay-disabled). + if let Some((raw_r, eff_r)) = retrievability { Self::update_record_state( - *memory_id, phase, eff_r, &meta, storage, graph, result, + *memory_id, phase, raw_r, eff_r, &meta, storage, graph, result, ) .await; } @@ -837,9 +865,9 @@ impl DecaySweepRunner { /// retrievability, and pushes a `PendingTransition` if the /// effective R falls below the phase threshold. /// - /// Returns `Some(effective_r)` if the record was evaluated (not - /// skipped), or `None` if the record was skipped (permastore or - /// decay disabled). + /// Returns `Some((raw_r, effective_r))` if the record was evaluated + /// (not skipped), or `None` if the record was skipped (permastore + /// or decay disabled). #[allow(clippy::too_many_arguments)] async fn evaluate_record( memory_id: MemoryId, @@ -852,7 +880,7 @@ impl DecaySweepRunner { global_decay_multiplier: f64, pending_transitions: &mut Vec, result: &mut SweepResult, - ) -> Option { + ) -> Option<(f32, f32)> { // -- Permastore exemption --------------------------------- if meta.is_permastore { result.permastore_skipped += 1; @@ -896,6 +924,7 @@ impl DecaySweepRunner { pending_transitions.push(PendingTransition { memory_id, from_phase: phase, + raw_r, effective_r, stability: meta.stability, is_permastore: meta.is_permastore, @@ -907,21 +936,27 @@ impl DecaySweepRunner { result.saved_by_connection_bonus += 1; } - Some(effective_r) + Some((raw_r, effective_r)) } - /// Update a record's decay_strength in storage and sync the - /// graph node state after evaluation. + /// Update a record's strength and decay_strength in storage and + /// sync the graph node state after evaluation. + /// + /// `raw_r` is the FSRS retrievability without connection bonus + /// (written to `record.strength`). + /// `effective_r` is the retrievability with connection bonus + /// (written to `record.decay_strength` and synced to graph). async fn update_record_state( memory_id: MemoryId, phase: DecayPhase, + raw_r: f32, effective_r: f32, meta: &DecayMetadata, storage: &Arc>, graph: &SharedGraph, result: &mut SweepResult, ) { - // Update decay_strength in metadata regardless of transition. + // Update strength and decay_strength in metadata regardless of transition. { let storage_clone = Arc::clone(storage); let stability = meta.stability; @@ -934,6 +969,7 @@ impl DecaySweepRunner { storage_r.update_decay_state( memory_id, phase, + raw_r, effective_r, stability, is_permastore, @@ -945,14 +981,14 @@ impl DecaySweepRunner { result.errors.push(SweepRecordError { memory_id, phase, - error: format!("failed to update decay_strength: {e}"), + error: format!("failed to update decay state: {e}"), }); } Err(e) => { result.errors.push(SweepRecordError { memory_id, phase, - error: format!("update decay_strength task panicked: {e}"), + error: format!("update decay state task panicked: {e}"), }); } _ => {} @@ -1000,6 +1036,7 @@ impl DecaySweepRunner { DecayPhase::Full => { Self::transition_full_to_summary( t.memory_id, + t.raw_r, t.effective_r, t.stability, t.is_permastore, @@ -1011,6 +1048,7 @@ impl DecaySweepRunner { DecayPhase::Summary => { Self::transition_summary_to_ghost( t.memory_id, + t.raw_r, t.effective_r, t.stability, t.is_permastore, @@ -1070,6 +1108,7 @@ impl DecaySweepRunner { /// the scan-time metadata snapshot so we avoid re-reading the record. async fn transition_full_to_summary( memory_id: MemoryId, + raw_r: f32, effective_r: f32, stability: f32, is_permastore: bool, @@ -1089,14 +1128,23 @@ impl DecaySweepRunner { let storage_r = storage_clone .read() .map_err(|e| SweepError::Storage(format!("storage lock poisoned: {e}")))?; + + // Update decay state: phase -> Summary, with raw and effective R. storage_r .update_decay_state( memory_id, DecayPhase::Summary, + raw_r, effective_r, stability, is_permastore, ) + .map_err(|e| SweepError::Storage(e.to_string()))?; + + // Zero the full_text pointer so the text data becomes dead + // space reclaimable by compaction. + storage_r + .strip_full_text(memory_id) .map_err(|e| SweepError::Storage(e.to_string())) }) .await @@ -1122,6 +1170,7 @@ impl DecaySweepRunner { /// metadata snapshot so we avoid re-reading the record. async fn transition_summary_to_ghost( memory_id: MemoryId, + raw_r: f32, effective_r: f32, stability: f32, is_permastore: bool, @@ -1132,14 +1181,23 @@ impl DecaySweepRunner { let storage_r = storage_clone .read() .map_err(|e| SweepError::Storage(format!("storage lock poisoned: {e}")))?; + + // Update decay state: phase -> Ghost, with raw and effective R. storage_r .update_decay_state( memory_id, DecayPhase::Ghost, + raw_r, effective_r, stability, is_permastore, ) + .map_err(|e| SweepError::Storage(e.to_string()))?; + + // Clear the summary field -- Ghost phase retains only + // the embedding and relationship edges. + storage_r + .strip_summary(memory_id) .map_err(|e| SweepError::Storage(e.to_string())) }) .await @@ -1172,22 +1230,52 @@ impl DecaySweepRunner { graph: &SharedGraph, ) -> Result { // 1. Remove all edges via the graph's bridging-aware deletion. - // Acquires a WRITE lock on the graph. + // Acquires a WRITE lock on the graph. If a bridge edge is + // created, resolve its NodeKey endpoints to MemoryIds while + // the lock is held (NodeKeys are not valid after release). + let bridge_to_persist: Option; { let mut graph_w = graph.write().await; let removal_result = graph_w.remove_memory_with_bridging(memory_id); debug!( id = %memory_id, edges_removed = removal_result.removed_edges.len(), + bridge_created = removal_result.bridge_created.is_some(), "removed edges with bridging" ); + + // Resolve the bridge's NodeKey endpoints to MemoryIds + // while the graph lock is still held. + bridge_to_persist = removal_result.bridge_created.and_then(|bridge| { + let source_id = graph_w.get_node_by_key(bridge.source).map(|n| n.memory_id); + let target_id = graph_w.get_node_by_key(bridge.target).map(|n| n.memory_id); + match (source_id, target_id) { + (Some(src), Some(tgt)) => Some(crate::storage::edges::PersistedEdge { + source: src, + target: tgt, + edge_type: bridge.edge_type, + weight: bridge.weight, + auto_created: bridge.auto_created, + created_at: bridge.created_at as u64, + }), + _ => { + warn!( + id = %memory_id, + "bridge edge created but source/target nodes \ + not found in graph; skipping persistence" + ); + None + } + } + }); } // write lock released - // 2. Remove edge records and metadata via spawn_blocking. + // 2. Remove old edge records, persist bridge, and delete + // metadata via spawn_blocking. { let storage_clone = Arc::clone(storage); tokio::task::spawn_blocking(move || { - // Remove edges. + // Remove edges for the deleted memory. { let storage_r = storage_clone .read() @@ -1195,6 +1283,20 @@ impl DecaySweepRunner { storage_r .remove_all_edges(memory_id) .map_err(|e| SweepError::Storage(e.to_string()))?; + + // Persist the bridge edge if one was created. + if let Some(ref bridge) = bridge_to_persist { + storage_r + .batch_add_edges(&[bridge.clone()]) + .map_err(|e| SweepError::Storage(e.to_string()))?; + tracing::debug!( + source = %bridge.source, + target = %bridge.target, + edge_type = ?bridge.edge_type, + weight = bridge.weight, + "persisted bridge edge to edges.db" + ); + } } // Delete metadata record (requires write lock). { diff --git a/src/embedding/bedrock.rs b/src/embedding/bedrock.rs index a86d33a..bf034e4 100644 --- a/src/embedding/bedrock.rs +++ b/src/embedding/bedrock.rs @@ -255,6 +255,15 @@ impl BedrockProvider { let response: CohereResponse = serde_json::from_slice(&response_bytes) .map_err(|e| EmbeddingError::InvalidResponse(e.to_string()))?; + // Validate the number of returned embeddings matches the input count. + if response.embeddings.len() != texts.len() { + return Err(EmbeddingError::InvalidResponse(format!( + "expected {} embeddings, got {}", + texts.len(), + response.embeddings.len() + ))); + } + for embedding in &response.embeddings { if embedding.len() != self.dim { return Err(EmbeddingError::DimensionMismatch { diff --git a/src/embedding/cache.rs b/src/embedding/cache.rs index b8af54b..125bb65 100644 --- a/src/embedding/cache.rs +++ b/src/embedding/cache.rs @@ -140,6 +140,15 @@ impl EmbeddingProvider for CachedProvider { if !miss_texts.is_empty() { let miss_embeddings = self.inner.embed_batch(&miss_texts).await?; + // Validate the provider returned exactly one embedding per input. + if miss_embeddings.len() != miss_texts.len() { + return Err(EmbeddingError::InvalidResponse(format!( + "provider returned {} embeddings for {} inputs", + miss_embeddings.len(), + miss_texts.len() + ))); + } + // Step 4: Insert into cache and place into results. for (j, embedding) in miss_embeddings.into_iter().enumerate() { let original_index = miss_indices[j]; @@ -149,7 +158,8 @@ impl EmbeddingProvider for CachedProvider { } } - // Step 5: Unwrap all results (every slot is now Some). + // Step 5: Unwrap all results (every slot is now Some after + // cache hits + validated miss embeddings filled all positions). Ok(results.into_iter().map(|r| r.unwrap()).collect()) } diff --git a/src/embedding/ollama.rs b/src/embedding/ollama.rs index 6d589f9..1331854 100644 --- a/src/embedding/ollama.rs +++ b/src/embedding/ollama.rs @@ -70,22 +70,6 @@ impl OllamaProvider { Ok(()) } - - /// List available models on the Ollama instance. - pub async fn list_models(&self) -> Result, EmbeddingError> { - let response = self - .client - .get(format!("{}/api/tags", self.base_url)) - .send() - .await?; - - let tags: TagsResponse = response - .json() - .await - .map_err(|e| EmbeddingError::InvalidResponse(e.to_string()))?; - - Ok(tags.models.into_iter().map(|m| m.name).collect()) - } } #[async_trait::async_trait] @@ -175,6 +159,15 @@ impl EmbeddingProvider for OllamaProvider { "Ollama embedding request succeeded" ); + // Validate the number of returned embeddings matches the input count. + if body.embeddings.len() != texts.len() { + return Err(EmbeddingError::InvalidResponse(format!( + "expected {} embeddings, got {}", + texts.len(), + body.embeddings.len() + ))); + } + // Validate every returned vector's dimensionality. for embedding in &body.embeddings { if embedding.len() != self.dim { @@ -218,26 +211,3 @@ struct EmbedResponse { struct OllamaErrorResponse { error: String, } - -#[derive(Deserialize)] -struct TagsResponse { - models: Vec, -} - -#[derive(Deserialize)] -struct ModelInfo { - name: String, - #[serde(default)] - #[allow(dead_code)] - details: ModelDetails, -} - -#[derive(Default, Deserialize)] -struct ModelDetails { - #[serde(default)] - #[allow(dead_code)] - family: String, - #[serde(default)] - #[allow(dead_code)] - parameter_size: String, -} diff --git a/src/embedding/openai.rs b/src/embedding/openai.rs index aaeeb84..a15dce1 100644 --- a/src/embedding/openai.rs +++ b/src/embedding/openai.rs @@ -84,11 +84,13 @@ impl OpenAIProvider { /// Retries on: /// - 429 (rate limit): up to MAX_RETRIES, exponential backoff + jitter /// - 500/502/503 (server error): up to MAX_RETRIES, exponential backoff + /// - Network errors (connect failures, timeouts): up to MAX_RETRIES /// /// Does NOT retry on: /// - 400 (bad request): returns RequestFailed immediately /// - 401/403 (auth): returns Unavailable immediately /// - 404 (model not found): returns ModelNotFound immediately + /// - Non-transient network errors (e.g., invalid URL): returns immediately async fn request_with_retry( &self, texts: &[&str], @@ -104,14 +106,32 @@ impl OpenAIProvider { const MAX_RETRIES: u32 = 5; for attempt in 0..=MAX_RETRIES { - let response = self + let response = match self .client .post(format!("{}/v1/embeddings", self.base_url)) .header("Authorization", format!("Bearer {}", self.api_key.expose())) .header("Content-Type", "application/json") .json(&request_body) .send() - .await?; + .await + { + Ok(resp) => resp, + Err(e) if e.is_connect() || e.is_timeout() => { + if attempt == MAX_RETRIES { + return Err(EmbeddingError::Network(e)); + } + warn!( + model = %self.model, + attempt = attempt + 1, + %e, + "OpenAI connection failed, retrying" + ); + tokio::time::sleep(delay).await; + delay = std::cmp::min(delay * 2, Duration::from_secs(60)); + continue; + } + Err(e) => return Err(EmbeddingError::Network(e)), + }; let status = response.status(); @@ -241,6 +261,15 @@ impl EmbeddingProvider for OpenAIProvider { let mut data = response.data; data.sort_by_key(|d| d.index); + // Validate the number of returned embeddings matches the chunk size. + if data.len() != chunk.len() { + return Err(EmbeddingError::InvalidResponse(format!( + "expected {} embeddings, got {}", + chunk.len(), + data.len() + ))); + } + for item in data { if item.embedding.len() != self.dim { return Err(EmbeddingError::DimensionMismatch { diff --git a/src/embedding/prefix.rs b/src/embedding/prefix.rs index 4a8104e..e5648da 100644 --- a/src/embedding/prefix.rs +++ b/src/embedding/prefix.rs @@ -79,10 +79,10 @@ impl EmbeddingProvider for PrefixedProvider { async fn embed_query(&self, text: &str) -> Result, EmbeddingError> { if self.query_prefix.is_empty() { - self.inner.embed(text).await + self.inner.embed_query(text).await } else { let prefixed = Self::prefixed(&self.query_prefix, text); - self.inner.embed(&prefixed).await + self.inner.embed_query(&prefixed).await } } @@ -263,4 +263,91 @@ mod tests { let p2 = PrefixedProvider::new(Box::new(mock2), "".to_string(), "".to_string()); assert!(!p2.has_query_prefix()); } + + /// Mock provider that distinguishes embed() from embed_query() + /// by recording which method was called. + struct AsymmetricMockProvider { + calls: Arc>>, + dims: usize, + } + + impl AsymmetricMockProvider { + fn new(dims: usize) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + calls: calls.clone(), + dims, + }, + calls, + ) + } + } + + #[async_trait::async_trait] + impl EmbeddingProvider for AsymmetricMockProvider { + async fn embed(&self, text: &str) -> Result, EmbeddingError> { + self.calls.lock().unwrap().push(("embed", text.to_string())); + Ok(vec![1.0; self.dims]) + } + + async fn embed_query(&self, text: &str) -> Result, EmbeddingError> { + self.calls + .lock() + .unwrap() + .push(("embed_query", text.to_string())); + Ok(vec![2.0; self.dims]) + } + + async fn embed_batch(&self, texts: &[&str]) -> Result>, EmbeddingError> { + let mut calls = self.calls.lock().unwrap(); + for text in texts { + calls.push(("embed_batch", text.to_string())); + } + Ok(texts.iter().map(|_| vec![1.0; self.dims]).collect()) + } + + fn dimensions(&self) -> usize { + self.dims + } + + fn model_name(&self) -> &str { + "asymmetric-mock" + } + } + + #[tokio::test] + async fn test_embed_query_delegates_to_inner_embed_query() { + let (mock, calls) = AsymmetricMockProvider::new(4); + let provider = + PrefixedProvider::new(Box::new(mock), "doc: ".to_string(), "query: ".to_string()); + provider.embed_query("hello").await.unwrap(); + let calls = calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "embed_query"); + assert_eq!(calls[0].1, "query: hello"); + } + + #[tokio::test] + async fn test_embed_delegates_to_inner_embed() { + let (mock, calls) = AsymmetricMockProvider::new(4); + let provider = + PrefixedProvider::new(Box::new(mock), "doc: ".to_string(), "query: ".to_string()); + provider.embed("hello").await.unwrap(); + let calls = calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "embed"); + assert_eq!(calls[0].1, "doc: hello"); + } + + #[tokio::test] + async fn test_embed_query_empty_prefix_delegates_to_inner_embed_query() { + let (mock, calls) = AsymmetricMockProvider::new(4); + let provider = PrefixedProvider::new(Box::new(mock), "doc: ".to_string(), "".to_string()); + provider.embed_query("hello").await.unwrap(); + let calls = calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "embed_query"); + assert_eq!(calls[0].1, "hello"); + } } diff --git a/src/error.rs b/src/error.rs index 0aa762a..0235a4a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -103,20 +103,3 @@ pub enum RecalldError { /// Convenience alias used at the system level. pub type Result = std::result::Result; - -impl RecalldError { - /// Map this error to an HTTP status code for the API response. - pub fn status_code(&self) -> u16 { - match self { - Self::Config(_) => 400, - Self::NotReady { .. } => 503, - Self::Timeout { .. } => 503, - _ => 500, - } - } - - /// Determine whether the caller should retry the operation. - pub fn is_retryable(&self) -> bool { - matches!(self, Self::Timeout { .. } | Self::NotReady { .. }) - } -} diff --git a/src/graph/activation.rs b/src/graph/activation.rs index 9e216cb..39d17c3 100644 --- a/src/graph/activation.rs +++ b/src/graph/activation.rs @@ -213,6 +213,11 @@ pub fn connection_bonus( // Optional 2-hop with decay if config.include_2hop { + // Collect 1-hop keys so we can skip them at hop-2 (already counted) + let hop1_keys: HashSet = hop1.iter().map(|(k, _)| *k).collect(); + // Track best contribution per 2-hop node to dedup across paths + let mut hop2_best: HashMap = HashMap::new(); + for (neighbor_key, _) in &hop1 { if let Some(neighbor_node) = graph.nodes.get(*neighbor_key) { let hop2 = collect_neighbor_contributions( @@ -226,10 +231,20 @@ pub fn connection_bonus( if hop2_key == node_key { continue; } - bonus += contribution * config.scale * HOP_2_DECAY; + // Skip if hop-2 was already counted as a 1-hop neighbor + if hop1_keys.contains(&hop2_key) { + continue; + } + // Dedup: keep the max contribution across different 1-hop paths + let entry = hop2_best.entry(hop2_key).or_insert(0.0_f32); + *entry = entry.max(contribution); } } } + + for (_, contribution) in &hop2_best { + bonus += contribution * config.scale * HOP_2_DECAY; + } } bonus.min(MAX_CONNECTION_BONUS) diff --git a/src/graph/autolink.rs b/src/graph/autolink.rs index 4ec2e9e..36b5749 100644 --- a/src/graph/autolink.rs +++ b/src/graph/autolink.rs @@ -30,41 +30,6 @@ pub const DEFAULT_MAX_LINKS: usize = 15; /// Applied per-pair, capped at 0.05 regardless of shared tag count. pub const TAG_THRESHOLD_ADJUSTMENT: f32 = 0.05; -// ═══════════════════════════════════════════════════════════════════════ -// AutoLinkThresholds -// ═══════════════════════════════════════════════════════════════════════ - -/// Default similarity thresholds by embedding model family. -/// -/// These are starting points -- the actual threshold is stored in -/// namespace configuration and tunable at runtime. -/// -/// Thresholds are NOT portable across models. Score distributions -/// vary significantly (ada-002 compresses scores upward; MiniLM -/// spreads them wide). See Spec 05 section 3.2. -pub struct AutoLinkThresholds; - -impl AutoLinkThresholds { - /// Return the recommended threshold for a model family. - /// Falls back to 0.60 for unknown models. - pub fn default_for_model(model_name: &str) -> f32 { - let lower = model_name.to_lowercase(); - if lower.contains("ada-002") { - 0.79 - } else if lower.contains("text-embedding-3-small") - || lower.contains("text-embedding-3-large") - { - 0.50 - } else if lower.contains("minilm") || lower.contains("all-minilm") { - 0.60 - } else if lower.contains("nomic-embed") { - 0.60 - } else { - 0.60 // conservative default - } - } -} - // ═══════════════════════════════════════════════════════════════════════ // AutoLinkCandidate // ═══════════════════════════════════════════════════════════════════════ @@ -194,12 +159,18 @@ pub enum AutoLinkError { // persist_edges — shared helper for edge persistence + graph insertion // ═══════════════════════════════════════════════════════════════════════ -/// Persist a batch of edges to storage, insert them into the in-memory -/// graph, and update the edge count on the source memory. +/// Insert edges into the in-memory graph, persist accepted edges to +/// storage, and update the edge count on the source memory. /// /// This helper consolidates the boilerplate that was previously duplicated /// across `perform_autolink`, `perform_entity_link`, and `perform_temporal_link`. /// +/// Edges are validated in the graph FIRST, and only edges that the graph +/// accepts are persisted to storage. This prevents the TOCTOU race where +/// edges written to storage before graph validation could be rejected by +/// the graph but survive in edges.db, getting loaded back on restart +/// without the bidirectional duplicate check. +/// /// # Returns /// /// The number of edges successfully inserted into the graph. @@ -214,39 +185,45 @@ pub(crate) async fn persist_edges( return Ok(0); } - // Step 1: Persist edges to edges.db. + // Step 1: Write-lock graph for edge insertion. Validate edges in the + // graph first, collecting only the ones that succeed. + let accepted_edges: Vec = { + let mut graph_w = graph.write().await; + let mut accepted = Vec::with_capacity(persisted_edges.len()); + for pe in persisted_edges { + match graph_w.add_edge(pe.source, pe.target, pe.edge_type, pe.weight, true) { + Ok(_) => accepted.push(pe.clone()), + Err(crate::graph::GraphError::EdgeExists(_, _)) => {} + Err(crate::graph::GraphError::MemoryNotFound(_)) => {} + Err(e) => return Err(AutoLinkError::Graph(e)), + } + } + accepted + }; + + let edges_created = accepted_edges.len(); + + if edges_created == 0 { + return Ok(0); + } + + // Step 2: Persist only accepted edges to edges.db. { let storage = storage.clone(); - let edges_owned = persisted_edges.to_vec(); tokio::task::spawn_blocking(move || { let storage_r = storage .read() .map_err(|e| AutoLinkError::LockPoisoned(format!("storage lock poisoned: {e}")))?; storage_r - .batch_add_edges(&edges_owned) + .batch_add_edges(&accepted_edges) .map_err(|e| AutoLinkError::Storage(e.to_string())) }) .await .map_err(|e| AutoLinkError::Storage(format!("blocking task join error: {e}")))??; } - // Step 2: Write-lock graph for edge insertion. - let edges_created = { - let mut graph_w = graph.write().await; - let mut created = 0usize; - for pe in persisted_edges { - match graph_w.add_edge(pe.source, pe.target, pe.edge_type, pe.weight, true) { - Ok(_) => created += 1, - Err(crate::graph::GraphError::EdgeExists(_, _)) => {} - Err(crate::graph::GraphError::MemoryNotFound(_)) => {} - Err(e) => return Err(AutoLinkError::Graph(e)), - } - } - created - }; - // Step 3: Additive edge_count update. - if edges_created > 0 { + { let current_count = cache .get(new_memory_id) .await @@ -255,7 +232,7 @@ pub(crate) async fn persist_edges( let new_count = current_count + edges_created as u16; { let storage = storage.clone(); - let _ = tokio::task::spawn_blocking(move || { + match tokio::task::spawn_blocking(move || { let storage_r = storage.read().map_err(|e| { AutoLinkError::LockPoisoned(format!("storage lock poisoned: {e}")) })?; @@ -263,7 +240,24 @@ pub(crate) async fn persist_edges( .update_edge_count(new_memory_id, new_count) .map_err(|e| AutoLinkError::Storage(e.to_string())) }) - .await; + .await + { + Ok(Ok(())) => {} + Ok(Err(e)) => { + tracing::warn!( + memory_id = %new_memory_id, + error = %e, + "Failed to update edge count in storage; cache is correct but on-disk record may be stale" + ); + } + Err(e) => { + tracing::warn!( + memory_id = %new_memory_id, + error = %e, + "Edge count update task panicked; cache is correct but on-disk record may be stale" + ); + } + } } cache.update_edge_count(new_memory_id, new_count).await; } diff --git a/src/graph/mod.rs b/src/graph/mod.rs index bf12cd5..a857b90 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -37,8 +37,8 @@ pub use crate::storage::PersistedEdge; // Re-export CS-11 auto-link types and functions pub use autolink::{ - AutoLinkCandidate, AutoLinkError, AutoLinkThresholds, DEFAULT_MAX_LINKS, THRESHOLD_HARD_FLOOR, - auto_link, perform_autolink, perform_entity_link, perform_temporal_link, + AutoLinkCandidate, AutoLinkError, DEFAULT_MAX_LINKS, THRESHOLD_HARD_FLOOR, auto_link, + perform_autolink, perform_entity_link, perform_temporal_link, }; use std::sync::Arc; diff --git a/src/graph/structure.rs b/src/graph/structure.rs index c1cb7b6..75a3e55 100644 --- a/src/graph/structure.rs +++ b/src/graph/structure.rs @@ -309,52 +309,6 @@ impl RelationshipGraph { Ok(key) } - /// Remove a node and ALL its edges. Returns the removed edges - /// so the caller can delete them from edges.db. - /// - /// This is the simple removal path (no bridging). For ghost - /// deletion with pass-through bridging, use - /// `remove_memory_with_bridging`. - pub fn remove_node(&mut self, memory_id: MemoryId) -> Result, GraphError> { - let node_key = self.resolve_key(memory_id)?; - - let node = self - .nodes - .get(node_key) - .ok_or(GraphError::MemoryNotFound(memory_id))?; - - // Collect all edge keys before mutating - let edge_keys: Vec = node - .outgoing - .iter() - .chain(node.incoming.iter()) - .copied() - .collect(); - - let mut removed_edges = Vec::with_capacity(edge_keys.len()); - - for ek in edge_keys { - if let Some(edge) = self.edges.remove(ek) { - // Remove this edge key from the OTHER node's lists - let other_key = if edge.source == node_key { - edge.target - } else { - edge.source - }; - if let Some(other_node) = self.nodes.get_mut(other_key) { - other_node.outgoing.retain(|k| *k != ek); - other_node.incoming.retain(|k| *k != ek); - } - removed_edges.push(edge); - } - } - - self.nodes.remove(node_key); - self.id_index.remove(&memory_id); - - Ok(removed_edges) - } - /// Update a node's cached decay state. Called by decay sweeps /// after recomputing R(t,S) and checking phase transitions. pub fn update_node_state( @@ -670,8 +624,10 @@ impl RelationshipGraph { let predecessor = in_edge.source; let successor = out_edge.target; - // Only bridge if no edge already exists and weight is sufficient - let edge_exists = self.edge_exists_between_keys(predecessor, successor); + // Only bridge if no edge of the same type already exists (either direction) + // and weight is sufficient + let edge_exists = + self.has_typed_edge_between(predecessor, successor, in_edge.edge_type); if !edge_exists && bridge_weight > 0.1 { let bridge = GraphEdge { @@ -735,20 +691,6 @@ impl RelationshipGraph { bridge_created, } } - - /// Check if any edge exists from source to target (forward direction only). - fn edge_exists_between_keys(&self, source: NodeKey, target: NodeKey) -> bool { - if let Some(source_node) = self.nodes.get(source) { - for &ek in &source_node.outgoing { - if let Some(edge) = self.edges.get(ek) { - if edge.target == target { - return true; - } - } - } - } - false - } } // ── Internal Helpers ──────────────────────────────────────────────── diff --git a/src/health/report.rs b/src/health/report.rs index 806ae9e..bce6b3f 100644 --- a/src/health/report.rs +++ b/src/health/report.rs @@ -267,7 +267,13 @@ pub async fn compute_storage_breakdown( namespaces: &dyn NamespaceRegistry, namespace_filter: Option, ) -> StorageBreakdown { - let db_path = storage.storage_path(); + let db_path = match storage.storage_path() { + Ok(p) => p, + Err(e) => { + tracing::warn!(%e, "failed to read storage path for breakdown"); + return StorageBreakdown::default(); + } + }; let meta_db_bytes = file_size(&db_path.join("meta.db")).unwrap_or(0); let edges_db_bytes = file_size(&db_path.join("edges.db")).unwrap_or(0); diff --git a/src/main.rs b/src/main.rs index 6c823cf..12a0ea3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -421,18 +421,31 @@ async fn run_serve( system.query_engine().clone(), system.embedding().clone(), system.vector_index().clone(), + system.fts_index().clone(), + system.entity_index().clone(), )); let storage: Arc = Arc::new(StorageEngineAdapter::new( system.storage().clone(), system.cache().clone(), - system.embedding().clone(), )); let cache: Arc = Arc::new(RecordCacheAdapter::new(system.cache().clone())); let graph: Arc = - Arc::new(RelationshipGraphAdapter::new(system.graph().clone())); - let decay: Arc = - Arc::new(FsrsEngineAdapter::new(system.storage().clone(), true)); + Arc::new(RelationshipGraphAdapter::new( + system.graph().clone(), + system.vector_index().clone(), + system.entity_index().clone(), + system.storage().clone(), + system.cache().clone(), + Arc::new(system.config().clone()), + )); + let decay: Arc = Arc::new(FsrsEngineAdapter::new( + system.storage().clone(), + system.cache().clone(), + system.graph().clone(), + Arc::new(system.config().clone()), + true, + )); let namespaces: Arc = Arc::new(NamespaceRegistryAdapter::new(system.storage().clone())); let metrics: Arc = Arc::new(NoopMetricsCollector); diff --git a/src/mcp/bridge_adapters.rs b/src/mcp/bridge_adapters.rs index 15fbab6..0c3303c 100644 --- a/src/mcp/bridge_adapters.rs +++ b/src/mcp/bridge_adapters.rs @@ -1000,16 +1000,49 @@ impl bridge::StorageEngine for McpStorageAdapter { // preserve the graph node and edges so relationship chains // remain intact for spreading activation traversal. - // 1. Read the current record to get tags for entity index cleanup. + // 1. Read, check, tombstone, and free the vector slot atomically + // under a single write lock to prevent TOCTOU races. Without + // this, two concurrent deletes could both read the record as + // non-tombstoned, both tombstone it, and both free the same + // vector slot -- corrupting the free list. let existing_record = { let storage = self.storage.clone(); tokio::task::spawn_blocking(move || { - let storage_r = storage.read().map_err(|e| { + let mut storage_w = storage.write().map_err(|e| { bridge::BridgeError::Internal(format!("storage lock poisoned: {e}")) })?; - storage_r + + let existing = storage_w .get_record(id) - .map_err(|e| bridge::BridgeError::Storage(e.to_string())) + .map_err(|e| bridge::BridgeError::Storage(e.to_string()))?; + + let Some(existing_record) = existing else { + return Ok::, bridge::BridgeError>( + None, + ); + }; + + if existing_record.phase == DecayPhase::Tombstone { + return Ok(None); + } + + // Tombstone the record. + storage_w + .tombstone_memory(id) + .map_err(|e| bridge::BridgeError::Storage(e.to_string()))?; + + // Free the vector slot. + let ns_id = NamespaceId::new(existing_record.namespace_id); + if let Err(e) = storage_w.free_vector_slot(ns_id, existing_record.vector_slot) { + tracing::warn!( + memory_id = %id, + vector_slot = existing_record.vector_slot, + %e, + "vector slot free failed (non-fatal)" + ); + } + + Ok(Some(existing_record)) }) .await .map_err(|e| { @@ -1021,50 +1054,10 @@ impl bridge::StorageEngine for McpStorageAdapter { return Ok(false); }; - if existing_record.phase == DecayPhase::Tombstone { - return Ok(false); - } - - // 2. Tombstone the record and free the vector slot. - { - let storage = self.storage.clone(); - let vector_slot = existing_record.vector_slot; - let namespace_id = existing_record.namespace_id; - tokio::task::spawn_blocking(move || { - { - let storage_r = storage.read().map_err(|e| { - bridge::BridgeError::Internal(format!("storage lock poisoned: {e}")) - })?; - storage_r - .tombstone_memory(id) - .map_err(|e| bridge::BridgeError::Storage(e.to_string()))?; - } - { - let mut storage_w = storage.write().map_err(|e| { - bridge::BridgeError::Internal(format!("storage lock poisoned: {e}")) - })?; - let ns_id = NamespaceId::new(namespace_id); - if let Err(e) = storage_w.free_vector_slot(ns_id, vector_slot) { - tracing::warn!( - memory_id = %id, - vector_slot, - %e, - "vector slot free failed (non-fatal)" - ); - } - } - Ok::<_, bridge::BridgeError>(()) - }) - .await - .map_err(|e| { - bridge::BridgeError::Internal(format!("blocking task join error: {e}")) - })??; - } - - // 3. Invalidate cache entry. + // 2. Invalidate cache entry. self.cache.invalidate(id).await; - // 4. Remove from FTS5 index. + // 3. Remove from FTS5 index. { let fts = self.fts_index.lock().await; if let Err(e) = fts.remove(id) { @@ -1076,7 +1069,7 @@ impl bridge::StorageEngine for McpStorageAdapter { } } - // 5. Remove from vector index. + // 4. Remove from vector index. { use crate::search::VectorIndex; let mut vi = self.vector_index.write().await; @@ -1089,7 +1082,7 @@ impl bridge::StorageEngine for McpStorageAdapter { } } - // 6. Remove from entity index. + // 5. Remove from entity index. { let metadata = crate::model::parse_structured_tags(&existing_record.tags); if !metadata.entities.is_empty() { @@ -1098,7 +1091,7 @@ impl bridge::StorageEngine for McpStorageAdapter { } } - // 7. Update graph node phase to Tombstone (keep node and edges). + // 6. Update graph node phase to Tombstone (keep node and edges). { let mut graph_w = self.graph.write().await; let _ = graph_w.update_node_state(id, DecayPhase::Tombstone, 0.0); @@ -1139,27 +1132,46 @@ impl bridge::StorageEngine for McpStorageAdapter { let phase = record.phase; - // Step 2: Compute elapsed days since last access. + // Step 2: Look up the namespace's decay rate multiplier. + let ns_decay_multiplier = { + let storage = self.storage.clone(); + let ns_id = record.namespace_id; + let global_multiplier = self.config.decay.decay_rate_multiplier; + tokio::task::spawn_blocking(move || { + let storage_r = storage.read().map_err(|e| { + bridge::BridgeError::Internal(format!("storage lock poisoned: {e}")) + })?; + let multiplier = storage_r + .get_namespace(NamespaceId::new(ns_id)) + .ok() + .flatten() + .and_then(|ns| ns.decay_rate_multiplier) + .unwrap_or(global_multiplier as f32); + Ok::(multiplier) + }) + .await + .map_err(|e| { + bridge::BridgeError::Internal(format!("blocking task join error: {e}")) + })?? + }; + + // Step 3: Compute elapsed days since last access. let now = chrono::Utc::now().timestamp_millis(); let elapsed_millis = (now - record.last_accessed_at).max(0) as f64; let elapsed_days = (elapsed_millis / 86_400_000.0) as f32; - // Step 3: Use FSRS to compute new stability based on quality rating. + // Step 4: Use FSRS to compute new stability based on quality rating. let decay_config = crate::decay::DecayConfig::default(); let engine = crate::decay::FsrsEngine::new(&decay_config); - let new_stability = engine.review_stability( - record.stability, - elapsed_days, - quality, - 1.0, // default decay rate multiplier - ); - - // Step 4: Compute new retrievability (just reinforced = 1.0). + let new_stability = + engine.review_stability(record.stability, elapsed_days, quality, ns_decay_multiplier); + + // Step 5: Compute new retrievability (just reinforced = 1.0). let new_strength = 1.0_f32; let is_permastore = record.is_permastore != 0 || new_stability >= decay_config.permastore_threshold; - // Step 5: Persist the updated decay state and access event. + // Step 6: Persist the updated decay state and access event. { let storage = self.storage.clone(); tokio::task::spawn_blocking(move || { @@ -1167,7 +1179,14 @@ impl bridge::StorageEngine for McpStorageAdapter { bridge::BridgeError::Internal(format!("storage lock poisoned: {e}")) })?; storage_r - .update_decay_state(id, phase, new_strength, new_stability, is_permastore) + .update_decay_state( + id, + phase, + new_strength, + new_strength, + new_stability, + is_permastore, + ) .map_err(|e| bridge::BridgeError::Storage(e.to_string()))?; storage_r .update_access(id, now, crate::model::AccessKind::ManualReinforcement) @@ -1179,7 +1198,7 @@ impl bridge::StorageEngine for McpStorageAdapter { })??; } - // Step 6: Update cache. + // Step 7: Update cache. if let Some(existing) = self.cache.get(id).await { let mut updated = (*existing).clone(); updated.stability = new_stability; @@ -1349,7 +1368,7 @@ impl bridge::NamespaceRegistry for McpNamespaceAdapter { id: NamespaceId::UNSET, name: input_name.clone(), embedding_dim: dim, - initial_stability: initial_stability.unwrap_or(3.7), + initial_stability: initial_stability.unwrap_or(3.7145), default_difficulty: 5.0, phase_thresholds: crate::model::namespace::PhaseThresholds::default(), permastore_threshold: 1500.0, @@ -1387,25 +1406,67 @@ impl bridge::NamespaceRegistry for McpNamespaceAdapter { let storage_r = storage.read().map_err(|e| { bridge::BridgeError::Internal(format!("storage lock poisoned: {e}")) })?; - let _ns = storage_r + let ns = storage_r .get_namespace_by_name(&ns_name) .map_err(|e| bridge::BridgeError::Storage(e.to_string()))? .ok_or_else(|| { bridge::BridgeError::NotFound(format!("namespace '{ns_name}' not found")) })?; + // Compute real stats by scanning all records in the namespace. + let all_records = storage_r + .scan_all() + .map_err(|e| bridge::BridgeError::Storage(e.to_string()))?; + + let mut memory_count: u64 = 0; + let mut full_count: u64 = 0; + let mut summary_count: u64 = 0; + let mut ghost_count: u64 = 0; + let mut permastore_count: u64 = 0; + let mut strength_sum: f64 = 0.0; + let mut edge_count: u64 = 0; + + for (_id, record) in &all_records { + if NamespaceId::new(record.namespace_id) != ns.id { + continue; + } + memory_count += 1; + + match record.phase { + DecayPhase::Full => full_count += 1, + DecayPhase::Summary => summary_count += 1, + DecayPhase::Ghost => ghost_count += 1, + DecayPhase::Tombstone => {} + } + + if record.is_permastore != 0 { + permastore_count += 1; + } + + strength_sum += record.decay_strength as f64; + edge_count += record.edge_count as u64; + } + + let avg_strength = if memory_count > 0 { + (strength_sum / memory_count as f64) as f32 + } else { + 0.0 + }; + + let vector_bytes = memory_count * ns.embedding_dim as u64 * 4; + Ok(bridge::NamespaceStats { name: ns_name, - memory_count: 0, + memory_count, phase_counts: bridge::PhaseCounts { - full: 0, - summary: 0, - ghost: 0, + full: full_count, + summary: summary_count, + ghost: ghost_count, }, - permastore_count: 0, - avg_strength: 0.0, - edge_count: 0, - vector_bytes: 0, + permastore_count, + avg_strength, + edge_count, + vector_bytes, }) }) .await diff --git a/src/model/error.rs b/src/model/error.rs index bf094bf..dd30092 100644 --- a/src/model/error.rs +++ b/src/model/error.rs @@ -30,7 +30,7 @@ pub enum TagError { // ValidationError // ═══════════════════════════════════════════════════════════════════════ -/// Errors from `Memory::validate()` and `CreateMemory` validation. +/// Errors from `Memory::validate()` and memory creation validation. /// /// Each variant carries a machine-readable `code()` suitable for the /// JSON error response `"error"` field, plus a human-readable message diff --git a/src/model/memory.rs b/src/model/memory.rs index b282d9b..21bf2c0 100644 --- a/src/model/memory.rs +++ b/src/model/memory.rs @@ -1,5 +1,5 @@ -//! Memory record types: the public API representation, creation request, -//! access events, and access kind classification. +//! Memory record types: the public API representation, access events, +//! and access kind classification. use serde::{Deserialize, Serialize}; @@ -190,107 +190,3 @@ impl Memory { Ok(()) } } - -// ═══════════════════════════════════════════════════════════════════════ -// CreateMemory — API request body -// ═══════════════════════════════════════════════════════════════════════ - -/// Request body for creating a new memory. -/// -/// Tags arrive as raw strings and are validated into `Tag` values on the -/// server side. -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CreateMemory { - /// Target namespace name. Defaults to `"default"`. - #[serde(default = "default_namespace")] - pub namespace: String, - - /// Short description. Required, max 2000 bytes. - pub summary: String, - - /// Full content. Optional, max 1 MiB. - #[serde(default)] - pub full_text: Option, - - /// Raw tag strings. Validated into `Tag` values server-side. - #[serde(default)] - pub tags: Vec, - - /// Pre-computed embedding. If omitted, the server generates one. - #[serde(default)] - pub embedding: Option>, - - /// Explicit initial stability override (days). - /// Must be in [0.01, 36500.0]. Uses namespace default if omitted. - #[serde(default)] - pub initial_stability: Option, -} - -/// Default namespace name for `CreateMemory`. -fn default_namespace() -> String { - "default".to_string() -} - -impl CreateMemory { - /// Validate the request body. This checks content limits and tag - /// format only — namespace existence and embedding dimensions - /// require external context and are checked by the API layer. - pub fn validate(&self) -> Result<(), ValidationError> { - // Summary - if self.summary.is_empty() { - return Err(ValidationError::SummaryEmpty); - } - if self.summary.len() > SUMMARY_MAX_BYTES { - return Err(ValidationError::SummaryTooLong { - len: self.summary.len(), - max: SUMMARY_MAX_BYTES, - }); - } - - // Full text - if let Some(ref ft) = self.full_text { - if ft.len() > FULL_TEXT_MAX_BYTES { - return Err(ValidationError::FullTextTooLong { - len: ft.len(), - max: FULL_TEXT_MAX_BYTES, - }); - } - } - - // Tags: count - if self.tags.len() > MAX_TAGS { - return Err(ValidationError::TooManyTags { - count: self.tags.len(), - max: MAX_TAGS, - }); - } - - // Tags: format (attempt to construct each) - for raw in &self.tags { - Tag::new(raw.as_str()).map_err(ValidationError::from)?; - } - - // Embedding: finite check - if let Some(ref emb) = self.embedding { - for (i, &val) in emb.iter().enumerate() { - if !val.is_finite() { - return Err(ValidationError::NonFiniteEmbedding { index: i }); - } - } - } - - // Stability override - if let Some(s) = self.initial_stability { - if !(STABILITY_FLOOR..=STABILITY_CEILING).contains(&s) { - return Err(ValidationError::InvalidStability { - value: s, - min: STABILITY_FLOOR, - max: STABILITY_CEILING, - }); - } - } - - Ok(()) - } -} diff --git a/src/model/mod.rs b/src/model/mod.rs index 2750eb7..e03e9da 100644 --- a/src/model/mod.rs +++ b/src/model/mod.rs @@ -31,7 +31,7 @@ pub use self::decay::DecayPhase; pub use self::edge::EdgeType; pub use self::error::{DecodeError, TagError, ValidationError}; pub use self::id::{MemoryId, NamespaceId}; -pub use self::memory::{AccessEvent, AccessKind, CreateMemory, Memory}; +pub use self::memory::{AccessEvent, AccessKind, Memory}; pub use self::namespace::{NamespaceConfig, PhaseThresholds}; pub use self::record::{CachedRecord, DiskRecord}; pub use self::tag::{StructuredMetadata, Tag, entity_overlap, parse_structured_tags}; diff --git a/src/model/record.rs b/src/model/record.rs index 1621d7c..39f8218 100644 --- a/src/model/record.rs +++ b/src/model/record.rs @@ -477,31 +477,3 @@ impl From<&DiskRecord> for CachedRecord { } } } - -impl CachedRecord { - /// Hydrate into an API `Memory` struct for returning to callers. - /// - /// The `namespace_name` must be resolved by the caller from the - /// namespace registry. Fields not held in the cache (`full_text`, - /// `embedding`, `access_history`) are set to `None`. - pub fn to_memory(&self, namespace_name: String) -> Memory { - Memory { - id: self.id, - namespace: namespace_name, - created_at: self.created_at, - last_accessed_at: self.last_accessed_at, - summary: self.summary.clone(), - full_text: None, // loaded on demand - tags: self.tags.clone(), - phase: self.phase, - strength: self.strength, - decay_strength: self.decay_strength, - stability: self.stability, - difficulty: self.difficulty, - is_permastore: self.is_permastore, - edge_count: self.edge_count, - embedding: None, // lives in mmap'd vectors.dat - access_history: None, // loaded on demand - } - } -} diff --git a/src/model/tag.rs b/src/model/tag.rs index e13c074..1843e82 100644 --- a/src/model/tag.rs +++ b/src/model/tag.rs @@ -61,14 +61,6 @@ impl Tag { pub fn as_str(&self) -> &str { &self.0 } - - /// Wrap an already-validated, lowercased tag string without - /// re-running validation. Used by the binary decoder (CS-02) - /// where tags were validated at write time. **Not public** — - /// only the serialization layer should call this. - pub(crate) fn from_trusted(s: String) -> Self { - Tag(s) - } } impl Deref for Tag { diff --git a/src/search/adapters.rs b/src/search/adapters.rs index 1d0612e..fcbc783 100644 --- a/src/search/adapters.rs +++ b/src/search/adapters.rs @@ -267,6 +267,24 @@ impl MetadataStore for StorageMetadataAdapter { .await .map_err(|e| SearchError::Internal(format!("blocking task join error: {e}")))? } + + async fn scan_all(&self) -> Result> { + let storage = self.storage.clone(); + tokio::task::spawn_blocking(move || { + let storage_r = storage + .read() + .map_err(|e| SearchError::MetadataError(format!("storage lock poisoned: {e}")))?; + let all = storage_r + .scan_all() + .map_err(|e| SearchError::MetadataError(e.to_string()))?; + Ok(all + .iter() + .map(|(_id, disk)| CachedRecord::from(disk)) + .collect()) + }) + .await + .map_err(|e| SearchError::Internal(format!("blocking task join error: {e}")))? + } } // ── RifProcessor ─────────────────────────────────────────────────── @@ -486,10 +504,10 @@ impl GraphReader for SharedGraphReader { /// /// The EntityIndex is shared across all namespaces. The adapter /// accepts a `namespace_id` parameter but does not filter by it -/// in v1 because the entity index does not track which namespace -/// a memory belongs to. Namespace filtering happens later in the -/// pipeline when CachedRecords are loaded and the filter stage -/// discards records from other namespaces. +/// because the entity index does not track which namespace a +/// memory belongs to. Namespace filtering happens in the pipeline's +/// `passes_filters` stage, which checks each candidate's +/// `CachedRecord::namespace_id` against the query namespace. pub struct SharedEntityIndexReader { entity_index: Arc>, } diff --git a/src/search/index.rs b/src/search/index.rs index 7fc4f4c..c116833 100644 --- a/src/search/index.rs +++ b/src/search/index.rs @@ -352,12 +352,11 @@ impl VectorIndex for FlatVectorIndex { } // Extract results sorted by descending score. - // `into_sorted_vec()` returns ascending order (by min-heap's reversed Ord), - // and `.rev()` flips it to descending — no additional sort needed. + // `into_sorted_vec()` returns ascending order by the reversed Ord, + // which is already descending by actual score — no additional sort needed. let results: Vec = heap .into_sorted_vec() .into_iter() - .rev() .map(|se| VectorSearchResult { id: self.entries[se.index].id, score: se.score, diff --git a/src/search/pipeline.rs b/src/search/pipeline.rs index ea87740..a1ed063 100644 --- a/src/search/pipeline.rs +++ b/src/search/pipeline.rs @@ -91,6 +91,10 @@ pub trait MetadataStore: Send + Sync { async fn get(&self, id: &MemoryId) -> Result>; /// Batch load multiple records. Missing IDs are silently skipped. async fn get_batch(&self, ids: &[MemoryId]) -> Result>; + /// Scan all records in storage. Used by MetadataOnly queries. + async fn scan_all(&self) -> Result> { + Ok(Vec::new()) + } } /// RIF processor -- applies retrieval-induced forgetting. @@ -486,7 +490,7 @@ impl QueryEngine { // -- Stage 5: Apply Filters (entity overlap + metadata filters) --- let stage_start = Instant::now(); - self.apply_filters(&mut candidates, &query, &query_entities); + self.apply_filters(&mut candidates, &query, &query_entities, ns_config.id); timings.apply_filters_us = stage_start.elapsed().as_micros() as u64; // -- Stage 6: Calculate Effective R -------------------------------- @@ -662,6 +666,7 @@ impl QueryEngine { candidates: &mut Vec, query: &SearchQuery, query_entities: &[String], + namespace_id: NamespaceId, ) { // Entity overlap scoring. if !query_entities.is_empty() { @@ -671,8 +676,8 @@ impl QueryEngine { } } - // Metadata filters (ghost, min_score, tags, phase, strength). - candidates.retain(|c| Self::passes_filters(c, query)); + // Metadata filters (namespace, ghost, min_score, tags, phase, strength). + candidates.retain(|c| Self::passes_filters(c, query, namespace_id)); } /// Stage 7: Apply retrieval-induced forgetting suppressions. @@ -880,13 +885,17 @@ impl QueryEngine { let result_ids: Vec = filtered.iter().map(|s| s.memory_id).collect(); let records = self.load_records(&result_ids).await?; - let results: Vec = records + // Build a lookup map so we can iterate `filtered` (which preserves + // descending similarity order from the vector search) instead of + // `records` (whose order is scrambled by cache-hit / disk-load batching). + let record_map: std::collections::HashMap = + records.iter().map(|r| (r.id, r)).collect(); + + let results: Vec = filtered .iter() - .filter_map(|rec| { - let score = filtered - .iter() - .find(|s| s.memory_id == rec.id) - .map(|s| s.score); + .filter_map(|scored| { + let rec = record_map.get(&scored.memory_id)?; + let score = Some(scored.score); Some(SearchResult { memory_id: rec.id, created_at: rec.created_at, @@ -968,10 +977,7 @@ impl QueryEngine { filter: &SearchFilter, max_results: usize, ) -> Result> { - // Use get_batch with an empty slice to signal a full scan. - // The MetadataStore adapter can implement a storage-level scan - // path here; for now we return whatever the store provides. - let all_records = self.meta_store.get_batch(&[]).await?; + let all_records = self.meta_store.scan_all().await?; // Apply client-side filtering for namespace and query filters. let filtered: Vec = all_records @@ -1031,9 +1037,20 @@ impl QueryEngine { } /// Check if a candidate passes all query filters. - fn passes_filters(candidate: &Candidate, query: &SearchQuery) -> bool { + fn passes_filters( + candidate: &Candidate, + query: &SearchQuery, + namespace_id: NamespaceId, + ) -> bool { let record = &candidate.record; + // Namespace filter: discard candidates from other namespaces. + // Entity recall and graph expansion can inject cross-namespace + // candidates; this gate ensures they never reach the results. + if record.namespace_id != namespace_id { + return false; + } + // Tombstone filter: always exclude tombstoned memories from // search results. Their content has been stripped; they only // exist as graph relay nodes for spreading activation. diff --git a/src/serialization/json.rs b/src/serialization/json.rs index a8a691d..c0a2945 100644 --- a/src/serialization/json.rs +++ b/src/serialization/json.rs @@ -11,6 +11,7 @@ use serde::{Deserialize, Serialize}; use crate::model::decay::DecayPhase; use crate::model::id::MemoryId; use crate::model::memory::AccessEvent; +use crate::model::parse_structured_tags; use crate::model::record::CachedRecord; use crate::model::tag::Tag; @@ -45,45 +46,6 @@ pub struct ApiError { pub field: Option, } -/// Paginated response wrapper for list endpoints. -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct PaginatedResponse { - /// The page of results. - pub data: Vec, - /// Total number of records matching the query. - pub total: u64, - /// Offset into the full result set. - pub offset: u64, - /// Maximum number of results per page. - pub limit: u64, - - /// Server-side processing time in microseconds. - #[serde(skip_serializing_if = "Option::is_none")] - pub took_us: Option, -} - -/// Pagination query parameters (deserialized from query string). -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct PaginationParams { - /// Offset into the result set. Default: 0. - #[serde(default = "default_offset")] - pub offset: u64, - - /// Maximum results to return. Default: 50. - #[serde(default = "default_limit")] - pub limit: u64, -} - -fn default_offset() -> u64 { - 0 -} - -fn default_limit() -> u64 { - 50 -} - /// Response wrapper for search results, including relevance scores. #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] @@ -137,6 +99,15 @@ pub struct MemoryResponse { /// Validated tags attached to this memory. pub tags: Vec, + /// Named entities extracted from tags (entity/ prefix). + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub entities: Vec, + /// Topics extracted from tags (topic/ prefix). + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub topics: Vec, + /// Emotions extracted from tags (emotion/ prefix). + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub emotions: Vec, /// Current decay phase. pub phase: DecayPhase, /// Raw FSRS retrievability R, in [0.0, 1.0]. @@ -166,6 +137,7 @@ impl MemoryResponse { /// resolved namespace name. Embedding and access history /// are set to `None` -- the caller populates them if requested. pub fn from_cached(record: &CachedRecord, namespace_name: String) -> Self { + let metadata = parse_structured_tags(&record.tags); Self { id: record.id, namespace: namespace_name, @@ -174,6 +146,9 @@ impl MemoryResponse { summary: record.summary.clone(), full_text: None, // loaded on demand from fulltext.dat tags: record.tags.clone(), + entities: metadata.entities, + topics: metadata.topics, + emotions: metadata.emotions, phase: record.phase, strength: record.strength, decay_strength: record.decay_strength, @@ -224,25 +199,6 @@ fn default_namespace() -> String { "default".to_string() } -/// PATCH /memories/{id} -- Update mutable fields. -/// Only provided fields are updated; omitted fields are untouched. -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateMemoryRequest { - /// Replace the summary. - #[serde(default)] - pub summary: Option, - - /// Replace tags entirely (not additive). - #[serde(default)] - pub tags: Option>, - - /// Force a manual reinforcement (equivalent to an access - /// with AccessKind::ManualReinforcement). - #[serde(default)] - pub reinforce: Option, -} - /// POST /search -- Search for memories. #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] @@ -264,6 +220,18 @@ pub struct SearchRequest { #[serde(default)] pub tags: Vec, + /// Named entities to filter by (converted to entity/ tags). + #[serde(default)] + pub entities: Vec, + + /// Topics to filter by (converted to topic/ tags). + #[serde(default)] + pub topics: Vec, + + /// Emotions to filter by (converted to emotion/ tags). + #[serde(default)] + pub emotions: Vec, + /// Minimum decay strength to include in results. #[serde(default)] pub min_strength: Option, @@ -277,6 +245,14 @@ pub struct SearchRequest { #[serde(default)] pub depth: u8, + /// Optional inclusive lower bound for temporal filtering (millis since epoch). + #[serde(default)] + pub time_range_start: Option, + + /// Optional inclusive upper bound for temporal filtering (millis since epoch). + #[serde(default)] + pub time_range_end: Option, + /// Whether to include the embedding vectors in results. /// Default: false (saves bandwidth). #[serde(default)] @@ -300,7 +276,9 @@ pub struct NamespaceRequest { pub name: String, /// Embedding dimensionality. Fixed at creation time. - pub embedding_dim: u32, + /// When omitted, defaults to the same dimensions as the 'default' namespace. + #[serde(default)] + pub embedding_dim: Option, /// Initial stability for new memories in this namespace (days). #[serde(default = "default_initial_stability")] @@ -309,6 +287,12 @@ pub struct NamespaceRequest { /// Desired retention rate. Default: 0.9. #[serde(default = "default_desired_retention")] pub desired_retention: f32, + + /// Decay rate multiplier for this namespace. + /// None = inherit from global config. + /// Some(1.0) = normal, Some(0.0) = disabled, Some(2.0) = 2x slower decay. + #[serde(default)] + pub decay_rate_multiplier: Option, } fn default_initial_stability() -> f32 { diff --git a/src/serialization/mod.rs b/src/serialization/mod.rs index 81770a5..d3cf293 100644 --- a/src/serialization/mod.rs +++ b/src/serialization/mod.rs @@ -7,6 +7,5 @@ mod json; pub use json::{ ApiError, ApiResponse, CreateMemoryRequest, MemoryResponse, NamespaceRequest, - NamespaceResponse, PaginatedResponse, PaginationParams, SearchHit, SearchRequest, - SearchResponse, UpdateMemoryRequest, + NamespaceResponse, SearchHit, SearchRequest, SearchResponse, }; diff --git a/src/storage/edges.rs b/src/storage/edges.rs index d14fe94..779120b 100644 --- a/src/storage/edges.rs +++ b/src/storage/edges.rs @@ -226,33 +226,6 @@ impl EdgeStore { Ok(()) } - /// Remove a single edge from both forward and reverse indexes. - /// - /// Returns `Ok(true)` if the edge existed and was removed, - /// `Ok(false)` if the edge was not found. - pub fn remove_edge( - &self, - source: MemoryId, - target: MemoryId, - edge_type: EdgeType, - ) -> Result { - let fwd_key = encode_forward_key(&source, edge_type, &target); - let rev_key = encode_reverse_key(&target, edge_type, &source); - - let write_txn = self.db.begin_write()?; - let removed; - { - let mut fwd = write_txn.open_table(FORWARD_EDGES)?; - let mut rev = write_txn.open_table(REVERSE_EDGES)?; - - removed = fwd.remove(fwd_key.as_slice())?.is_some(); - rev.remove(rev_key.as_slice())?; - } - write_txn.commit()?; - - Ok(removed) - } - /// Return all outgoing edges from `source` (any edge type). /// /// Prefix scan on the first 16 bytes of the forward index key. @@ -309,37 +282,6 @@ impl EdgeStore { Ok(results) } - /// Return all outgoing edges of a specific type from `source`. - /// - /// Uses a 17-byte prefix: [source_uuid: 16][edge_type: 1]. - pub fn get_outgoing_of_type( - &self, - source: MemoryId, - edge_type: EdgeType, - ) -> Result, StorageError> { - let read_txn = self.db.begin_read()?; - let table = read_txn.open_table(FORWARD_EDGES)?; - - let mut start = [0u8; EDGE_KEY_SIZE]; - start[0..16].copy_from_slice(source.as_bytes()); - start[16] = edge_type.as_u8(); - - let mut end = [0u8; EDGE_KEY_SIZE]; - end[0..16].copy_from_slice(source.as_bytes()); - end[16] = edge_type.as_u8(); - end[17..33].fill(0xFF); - - let mut targets = Vec::new(); - for entry in table.range(start.as_slice()..=end.as_slice())? { - let (key_guard, _) = entry?; - let key = key_guard.value(); - let target = MemoryId::from_bytes(key[17..33].try_into().unwrap()); - targets.push(target); - } - - Ok(targets) - } - /// Remove all edges where `memory_id` appears as source or target. /// /// Returns the total number of edges removed. diff --git a/src/storage/engine.rs b/src/storage/engine.rs index 72abd56..ac4ba3c 100644 --- a/src/storage/engine.rs +++ b/src/storage/engine.rs @@ -15,6 +15,7 @@ use crate::model::{ }; use crate::storage::edges::{EdgeStore, PersistedEdge, cleanup_orphaned_edges}; use crate::storage::error::StorageError; +use crate::storage::fsync::{fsync_dir, fsync_file}; use crate::storage::metadata::MetadataStore; use crate::storage::text::{CompactionResult, TextRef, TextStore, recover_text_compaction}; use crate::storage::vectors::VectorManager; @@ -92,15 +93,32 @@ pub trait StorageEngine: Send + Sync { // ── Decay State ───────────────────────────────────────────────── /// Update the decay state of a single memory. + /// + /// `strength` is the raw FSRS retrievability R(t,S). + /// `decay_strength` is the effective retrievability including + /// connection bonus. fn update_decay_state( &self, id: MemoryId, phase: DecayPhase, strength: f32, + decay_strength: f32, stability: f32, is_permastore: bool, ) -> Result<(), StorageError>; + /// Zero the full_text pointer on a memory record. + /// + /// Called during Full -> Summary phase transitions to mark the + /// text data in fulltext.dat as dead space reclaimable by + /// compaction. + fn strip_full_text(&self, id: MemoryId) -> Result<(), StorageError>; + + /// Clear the summary field on a memory record. + /// + /// Called during Summary -> Ghost phase transitions. + fn strip_summary(&self, id: MemoryId) -> Result<(), StorageError>; + /// Return the IDs of all memories currently in the given phase. fn ids_in_phase(&self, phase: DecayPhase) -> Result, StorageError>; @@ -409,6 +427,19 @@ impl StorageEngine for RedbStorageEngine { record.text_length = text_ref.length; } + // 2b. Fsync vector and text data before the meta.db commit. + // meta.db (redb) fsyncs internally on commit, so without this + // step a crash could leave meta.db pointing to unflushed vector + // or text data. Vectors have no integrity check, so stale/zero + // data would be silently used for similarity search. Text has + // CRC protection but would still fail reads until rewritten. + if let Some(vs) = self.vector_manager.get(namespace_id) { + vs.sync()?; + } + if full_text.is_some() { + self.text_store.sync_data()?; + } + // 3. Insert the completed record into meta.db. if let Err(e) = self.meta_store.insert(id, record) { // Best-effort rollback: free the vector slot so it can be @@ -498,11 +529,26 @@ impl StorageEngine for RedbStorageEngine { id: MemoryId, phase: DecayPhase, strength: f32, + decay_strength: f32, stability: f32, is_permastore: bool, ) -> Result<(), StorageError> { - self.meta_store - .update_decay_state(id, phase, strength, stability, is_permastore) + self.meta_store.update_decay_state( + id, + phase, + strength, + decay_strength, + stability, + is_permastore, + ) + } + + fn strip_full_text(&self, id: MemoryId) -> Result<(), StorageError> { + self.meta_store.strip_full_text(id) + } + + fn strip_summary(&self, id: MemoryId) -> Result<(), StorageError> { + self.meta_store.strip_summary(id) } fn ids_in_phase(&self, phase: DecayPhase) -> Result, StorageError> { @@ -618,29 +664,60 @@ impl StorageEngine for RedbStorageEngine { } } - // Run compaction. - let result = self.text_store.compact(&live_refs)?; - - // Update meta.db with new text offsets. - let updates: Vec<(MemoryId, DiskRecord)> = result - .new_refs - .iter() - .map(|(id, new_ref)| { - // Find the original record and update its text pointer. - let original = all_records - .iter() - .find(|(rid, _)| rid == id) - .map(|(_, r)| r) - .expect("compaction ref must match a record"); - let mut updated = original.clone(); - updated.text_offset = new_ref.file_offset; - updated.text_length = new_ref.length; - (*id, updated) - }) - .collect(); - - if !updates.is_empty() { - self.meta_store.batch_update_records(&updates)?; + // Phase 1: Write new file with only live entries. + // This creates fulltext.dat.new and the .compacting marker, + // but does NOT rename or remove the marker. + let (result, compacted) = self.text_store.compact_write_new_file(&live_refs)?; + + if compacted { + // Phase 2: Update meta.db with new text offsets BEFORE + // renaming the file. This ensures that if we crash after + // the meta.db commit but before the rename, recovery can + // detect the .compaction-meta-committed marker and complete + // the rename. Without this ordering, a crash after the + // rename but before the meta.db update would leave stale + // offsets pointing into a compacted file, corrupting all + // text reads. + let updates: Vec<(MemoryId, DiskRecord)> = result + .new_refs + .iter() + .map(|(id, new_ref)| { + // Find the original record and update its text pointer. + let original = all_records + .iter() + .find(|(rid, _)| rid == id) + .map(|(_, r)| r) + .expect("compaction ref must match a record"); + let mut updated = original.clone(); + updated.text_offset = new_ref.file_offset; + updated.text_length = new_ref.length; + (*id, updated) + }) + .collect(); + + if !updates.is_empty() { + self.meta_store.batch_update_records(&updates)?; + } + + // Write a committed marker so recovery knows meta.db has + // already been updated with new offsets. If we crash after + // this but before the rename, recovery will complete the + // rename instead of rolling back. + let committed_marker = self.db_path.join(".compaction-meta-committed"); + std::fs::write(&committed_marker, b"meta.db offsets committed\n")?; + fsync_file(&committed_marker)?; + fsync_dir(&self.db_path)?; + + // Phase 3: Atomically swap the files and clean up markers. + self.text_store.compact_finalize()?; + + tracing::info!( + old_size = result.old_size, + new_size = result.new_size, + removed = result.entries_removed, + kept = result.entries_kept, + "Text.log compaction complete" + ); } Ok(result) diff --git a/src/storage/indexes.rs b/src/storage/indexes.rs index 9fbc698..a5e876a 100644 --- a/src/storage/indexes.rs +++ b/src/storage/indexes.rs @@ -1,28 +1,29 @@ //! Secondary index data structures for the metadata store. //! //! Contains `PhaseIndex` — roaring bitmap indexes tracking which vector -//! slots belong to each decay phase. Managed internally by -//! `MetadataStore`; persisted to `INDEX_TABLE` in meta.db. +//! slots belong to each decay phase, scoped per namespace to avoid +//! cross-namespace collisions. +//! +//! Managed internally by `MetadataStore`; persisted to `INDEX_TABLE` +//! in meta.db. //! //! See CS-07 Section 6. +use std::collections::HashMap; + use roaring::RoaringBitmap; use crate::model::DecayPhase; use crate::storage::error::StorageError; // ═══════════════════════════════════════════════════════════════════════ -// PhaseIndex +// NamespaceBitmaps — per-namespace set of phase bitmaps // ═══════════════════════════════════════════════════════════════════════ -/// In-memory bitmap index tracking which vector slots belong to each -/// decay phase. The authoritative copy lives here in RAM; it is -/// serialized to meta.db `INDEX_TABLE` periodically and rebuilt from -/// `META_TABLE` on startup if missing or corrupt. +/// Four roaring bitmaps (one per decay phase) for a single namespace. /// -/// Stores `vector_slot` (u32), not `MemoryId` (UUID is 128-bit and -/// does not fit roaring's u32 universe). -pub struct PhaseIndex { +/// Private to this module; external callers interact through `PhaseIndex`. +struct NamespaceBitmaps { /// Phase 1: full text + embedding. full: RoaringBitmap, /// Phase 2: summary + embedding. @@ -30,15 +31,11 @@ pub struct PhaseIndex { /// Phase 3: embedding only. ghost: RoaringBitmap, /// Phase 4: tombstone — content stripped, graph preserved. - /// Tombstoned memories are removed from the vector index, so this - /// bitmap is typically empty in practice (vector slots are freed). - /// It exists for completeness of the phase tracking. tombstone: RoaringBitmap, } -impl PhaseIndex { - /// Create an empty phase index. - pub fn new() -> Self { +impl NamespaceBitmaps { + fn new() -> Self { Self { full: RoaringBitmap::new(), summary: RoaringBitmap::new(), @@ -47,24 +44,7 @@ impl PhaseIndex { } } - /// Add a new memory (always starts in Phase::Full). - pub fn insert(&mut self, slot: u32) { - self.full.insert(slot); - } - - /// Move a memory between phases. - pub fn transition(&mut self, slot: u32, from: DecayPhase, to: DecayPhase) { - self.bitmap_mut(from).remove(slot); - self.bitmap_mut(to).insert(slot); - } - - /// Remove a memory entirely (after deletion). - pub fn remove(&mut self, slot: u32, phase: DecayPhase) { - self.bitmap_mut(phase).remove(slot); - } - - /// Get a read-only reference to a phase's bitmap. - pub fn bitmap(&self, phase: DecayPhase) -> &RoaringBitmap { + fn bitmap(&self, phase: DecayPhase) -> &RoaringBitmap { match phase { DecayPhase::Full => &self.full, DecayPhase::Summary => &self.summary, @@ -73,11 +53,6 @@ impl PhaseIndex { } } - /// Count of memories in a given phase. - pub fn count(&self, phase: DecayPhase) -> u64 { - self.bitmap(phase).len() - } - fn bitmap_mut(&mut self, phase: DecayPhase) -> &mut RoaringBitmap { match phase { DecayPhase::Full => &mut self.full, @@ -87,98 +62,253 @@ impl PhaseIndex { } } + fn is_empty(&self) -> bool { + self.full.is_empty() + && self.summary.is_empty() + && self.ghost.is_empty() + && self.tombstone.is_empty() + } +} + +// ═══════════════════════════════════════════════════════════════════════ +// PhaseIndex +// ═══════════════════════════════════════════════════════════════════════ + +/// In-memory bitmap index tracking which vector slots belong to each +/// decay phase, partitioned by namespace. +/// +/// The authoritative copy lives here in RAM; it is serialized to +/// meta.db `INDEX_TABLE` periodically and rebuilt from `META_TABLE` +/// on startup if missing or corrupt. +/// +/// Stores `(namespace_id, vector_slot)` pairs. Each namespace gets +/// its own set of four `RoaringBitmap`s (one per `DecayPhase`) to +/// avoid cross-namespace collisions — different namespaces allocate +/// vector slots independently and can share the same slot numbers. +pub struct PhaseIndex { + /// Bitmaps keyed by namespace_id (u32). + namespaces: HashMap, +} + +/// Serialization format version for namespace-aware bitmaps. +/// +/// The v1 format (pre-namespace) started with a u32 bitmap size that +/// was always >= 8 (the minimum serialized size of an empty +/// `RoaringBitmap`). Version 2 uses `2u32` as a distinguishing +/// marker in the first 4 bytes, which cannot collide with v1. +const FORMAT_VERSION: u32 = 2; + +impl PhaseIndex { + /// Create an empty phase index. + pub fn new() -> Self { + Self { + namespaces: HashMap::new(), + } + } + + /// Add a new memory (always starts in Phase::Full). + pub fn insert(&mut self, namespace_id: u32, slot: u32) { + self.namespaces + .entry(namespace_id) + .or_insert_with(NamespaceBitmaps::new) + .bitmap_mut(DecayPhase::Full) + .insert(slot); + } + + /// Move a memory between phases. + pub fn transition(&mut self, namespace_id: u32, slot: u32, from: DecayPhase, to: DecayPhase) { + let ns = self + .namespaces + .entry(namespace_id) + .or_insert_with(NamespaceBitmaps::new); + ns.bitmap_mut(from).remove(slot); + ns.bitmap_mut(to).insert(slot); + } + + /// Remove a memory entirely (after deletion). + pub fn remove(&mut self, namespace_id: u32, slot: u32, phase: DecayPhase) { + if let Some(ns) = self.namespaces.get_mut(&namespace_id) { + ns.bitmap_mut(phase).remove(slot); + } + } + + /// Return all `(namespace_id, vector_slot)` pairs in the given + /// phase across all namespaces. + pub fn all_slots_in_phase(&self, phase: DecayPhase) -> Vec<(u32, u32)> { + let mut result = Vec::new(); + for (&ns_id, ns_bitmaps) in &self.namespaces { + for slot in ns_bitmaps.bitmap(phase).iter() { + result.push((ns_id, slot)); + } + } + result + } + + /// Count of memories in a given phase (across all namespaces). + pub fn count(&self, phase: DecayPhase) -> u64 { + self.namespaces + .values() + .map(|ns| ns.bitmap(phase).len()) + .sum() + } + /// Serialize all bitmaps into a single byte buffer. /// - /// Format: `[u32 size_1][bitmap_1 bytes][u32 size_2][bitmap_2 bytes] - /// [u32 size_3][bitmap_3 bytes][u32 size_4][bitmap_4 bytes]` - /// Order: full, summary, ghost, tombstone. - /// - /// Backward-compatible: `from_bytes` tolerates the absence of the - /// tombstone bitmap (pre-tombstone data files). + /// Format v2 (namespace-aware): + /// ```text + /// [u32 version = 2] + /// [u32 namespace_count] + /// for each namespace (sorted by namespace_id for determinism): + /// [u32 namespace_id] + /// [u32 size_full] [full bitmap bytes] + /// [u32 size_summary] [summary bitmap bytes] + /// [u32 size_ghost] [ghost bitmap bytes] + /// [u32 size_tombstone] [tombstone bitmap bytes] + /// ``` pub fn to_bytes(&self) -> Vec { let mut buf = Vec::new(); - for bitmap in [&self.full, &self.summary, &self.ghost, &self.tombstone] { - let size = bitmap.serialized_size(); - buf.extend_from_slice(&(size as u32).to_le_bytes()); - bitmap - .serialize_into(&mut buf) - .expect("bitmap serialization is infallible"); + + // Version marker. + buf.extend_from_slice(&FORMAT_VERSION.to_le_bytes()); + + // Filter out empty namespaces to keep the serialized form compact. + let mut sorted_ns: Vec<(&u32, &NamespaceBitmaps)> = self + .namespaces + .iter() + .filter(|(_, ns)| !ns.is_empty()) + .collect(); + sorted_ns.sort_by_key(|(id, _)| **id); + + buf.extend_from_slice(&(sorted_ns.len() as u32).to_le_bytes()); + + for (ns_id, ns_bitmaps) in &sorted_ns { + buf.extend_from_slice(&ns_id.to_le_bytes()); + for bitmap in [ + &ns_bitmaps.full, + &ns_bitmaps.summary, + &ns_bitmaps.ghost, + &ns_bitmaps.tombstone, + ] { + let size = bitmap.serialized_size(); + buf.extend_from_slice(&(size as u32).to_le_bytes()); + bitmap + .serialize_into(&mut buf) + .expect("bitmap serialization is infallible"); + } } + buf } /// Deserialize from bytes produced by `to_bytes()`. /// - /// Backward-compatible: if only 3 bitmaps are present (pre-tombstone - /// format), the tombstone bitmap is initialized as empty. + /// Detects the format version: + /// - Version 2 (namespace-aware): parsed in full. + /// - Version 1 (pre-namespace, first u32 >= 8): returns an empty + /// `PhaseIndex`. The caller's startup validation will rebuild + /// from `META_TABLE` via `rebuild_from_records`. pub fn from_bytes(data: &[u8]) -> Result { - let mut offset = 0; + if data.len() < 4 { + return Err(StorageError::CorruptIndex( + "phase index data too short".into(), + )); + } + + let version = u32::from_le_bytes(data[0..4].try_into().unwrap()); + if version == FORMAT_VERSION { + Self::from_bytes_v2(data) + } else { + // Pre-namespace format. Return empty; rebuild_phase_index() + // will reconstruct correctly during startup validation. + Ok(Self::new()) + } + } + + /// Parse the v2 namespace-aware format. + fn from_bytes_v2(data: &[u8]) -> Result { + let mut offset = 4; // skip version marker + + if offset + 4 > data.len() { + return Err(StorageError::CorruptIndex( + "phase index v2 truncated at namespace count".into(), + )); + } + let ns_count = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()) as usize; + offset += 4; + + let mut namespaces = HashMap::with_capacity(ns_count); let labels = ["full", "summary", "ghost", "tombstone"]; - let mut bitmaps = Vec::with_capacity(4); - for (i, label) in labels.iter().enumerate() { - // The tombstone bitmap is optional for backward compatibility. - if offset >= data.len() { - if i >= 3 { - // Tombstone bitmap missing — use empty. - bitmaps.push(RoaringBitmap::new()); - break; - } - return Err(StorageError::CorruptIndex(format!( - "phase bitmap '{}' truncated at length prefix", - label - ))); - } + + for _ in 0..ns_count { if offset + 4 > data.len() { - if i >= 3 { - bitmaps.push(RoaringBitmap::new()); - break; - } - return Err(StorageError::CorruptIndex(format!( - "phase bitmap '{}' truncated at length prefix", - label - ))); + return Err(StorageError::CorruptIndex( + "phase index v2 truncated at namespace id".into(), + )); } - let size = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()) as usize; + let ns_id = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()); offset += 4; - if offset + size > data.len() { - return Err(StorageError::CorruptIndex(format!( - "phase bitmap '{}' truncated at payload", - label - ))); + + let mut bitmaps = Vec::with_capacity(4); + for label in &labels { + if offset + 4 > data.len() { + return Err(StorageError::CorruptIndex(format!( + "phase index v2 truncated at '{}' bitmap size for ns {}", + label, ns_id + ))); + } + let size = + u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()) as usize; + offset += 4; + + if offset + size > data.len() { + return Err(StorageError::CorruptIndex(format!( + "phase index v2 truncated at '{}' bitmap payload for ns {}", + label, ns_id + ))); + } + let bitmap = RoaringBitmap::deserialize_from(&data[offset..offset + size]) + .map_err(|e| { + StorageError::CorruptIndex(format!( + "invalid roaring bitmap for '{}' in ns {}: {}", + label, ns_id, e + )) + })?; + offset += size; + bitmaps.push(bitmap); } - let bitmap = - RoaringBitmap::deserialize_from(&data[offset..offset + size]).map_err(|e| { - StorageError::CorruptIndex(format!( - "invalid roaring bitmap for '{}': {}", - label, e - )) - })?; - offset += size; - bitmaps.push(bitmap); - } - // Ensure we have exactly 4 bitmaps. - while bitmaps.len() < 4 { - bitmaps.push(RoaringBitmap::new()); + + namespaces.insert( + ns_id, + NamespaceBitmaps { + full: bitmaps.remove(0), + summary: bitmaps.remove(0), + ghost: bitmaps.remove(0), + tombstone: bitmaps.remove(0), + }, + ); } - Ok(Self { - full: bitmaps.remove(0), - summary: bitmaps.remove(0), - ghost: bitmaps.remove(0), - tombstone: bitmaps.remove(0), - }) + + Ok(Self { namespaces }) } /// Rebuild from a full META_TABLE scan. Used on startup when the /// persisted bitmap is missing or corrupt. - pub fn rebuild_from_records(records: &[(u32, DecayPhase)]) -> Self { + /// + /// Each tuple is `(namespace_id, vector_slot, phase)`. + pub fn rebuild_from_records(records: &[(u32, u32, DecayPhase)]) -> Self { let mut index = Self::new(); - for &(slot, phase) in records { + for &(namespace_id, slot, phase) in records { // Tombstoned records have had their vector slot freed, so // they should not be tracked in the bitmap. Skip them. if phase == DecayPhase::Tombstone { continue; } - index.bitmap_mut(phase).insert(slot); + index + .namespaces + .entry(namespace_id) + .or_insert_with(NamespaceBitmaps::new) + .bitmap_mut(phase) + .insert(slot); } index } diff --git a/src/storage/metadata.rs b/src/storage/metadata.rs index d7cc6c5..cdb51a5 100644 --- a/src/storage/metadata.rs +++ b/src/storage/metadata.rs @@ -14,6 +14,7 @@ use redb::{ TableDefinition, }; +use crate::model::constants::ACCESS_HISTORY_MAX; use crate::model::{ AccessEvent, AccessKind, DecayPhase, DiskRecord, MemoryId, NamespaceConfig, NamespaceId, }; @@ -184,8 +185,8 @@ impl MetadataStore { // 4. Update in-memory phase bitmap (always Phase::Full for new // records). { - let mut pi = self.phase_index.write().unwrap(); - pi.insert(record.vector_slot); + let mut pi = self.phase_index.write().unwrap_or_else(|e| e.into_inner()); + pi.insert(record.namespace_id, record.vector_slot); } Ok(()) @@ -208,14 +209,20 @@ impl MetadataStore { /// Update the decay fields of a single memory record. /// - /// Read-modify-write: loads the existing record, patches the four - /// decay fields, re-serializes, and writes back. If the phase + /// Read-modify-write: loads the existing record, patches the five + /// decay fields (phase, strength, decay_strength, stability, + /// is_permastore), re-serializes, and writes back. If the phase /// changed, the phase bitmap index is updated. + /// + /// `strength` is the raw FSRS retrievability R(t,S). + /// `decay_strength` is the effective retrievability including + /// connection bonus. pub fn update_decay_state( &self, id: MemoryId, phase: DecayPhase, strength: f32, + decay_strength: f32, stability: f32, is_permastore: bool, ) -> Result<(), StorageError> { @@ -225,6 +232,7 @@ impl MetadataStore { old_phase = Some(record.phase); record.phase = phase; record.strength = strength; + record.decay_strength = decay_strength; record.stability = stability; record.is_permastore = if is_permastore { 1 } else { 0 }; })?; @@ -233,17 +241,41 @@ impl MetadataStore { // Update phase bitmap if phase changed. if old_phase != phase { - let mut pi = self.phase_index.write().unwrap(); - pi.transition(record.vector_slot, old_phase, phase); + let mut pi = self.phase_index.write().unwrap_or_else(|e| e.into_inner()); + pi.transition(record.namespace_id, record.vector_slot, old_phase, phase); } Ok(()) } + /// Zero the full_text pointer on a memory record. + /// + /// Called during Full -> Summary phase transitions to mark the + /// text data in fulltext.dat as dead space reclaimable by + /// compaction. Does not modify fulltext.dat itself. + pub fn strip_full_text(&self, id: MemoryId) -> Result<(), StorageError> { + self.update_record(&id, |record| { + record.text_offset = 0; + record.text_length = 0; + })?; + Ok(()) + } + + /// Clear the summary field on a memory record. + /// + /// Called during Summary -> Ghost phase transitions. After this, + /// only the embedding and relationship edges remain. + pub fn strip_summary(&self, id: MemoryId) -> Result<(), StorageError> { + self.update_record(&id, |record| { + record.summary = String::new(); + })?; + Ok(()) + } + /// Record a new access event on a memory. /// /// Updates `last_accessed_at` and appends to the access history - /// ring buffer (capped at 32 entries; drops oldest if full). + /// ring buffer (capped at ACCESS_HISTORY_MAX entries; drops oldest if full). pub fn update_access( &self, id: MemoryId, @@ -254,7 +286,7 @@ impl MetadataStore { record.last_accessed_at = timestamp; let event = AccessEvent { timestamp, kind }; - if record.access_history.len() >= 32 { + if record.access_history.len() >= ACCESS_HISTORY_MAX { record.access_history.remove(0); } record.access_history.push(event); @@ -302,8 +334,12 @@ impl MetadataStore { // 4. Remove from in-memory phase bitmap. { - let mut pi = self.phase_index.write().unwrap(); - pi.remove(deleted_record.vector_slot, deleted_record.phase); + let mut pi = self.phase_index.write().unwrap_or_else(|e| e.into_inner()); + pi.remove( + deleted_record.namespace_id, + deleted_record.vector_slot, + deleted_record.phase, + ); } Ok(Some(deleted_record)) @@ -327,6 +363,7 @@ impl MetadataStore { let old_phase; let old_tags; let old_vector_slot; + let old_namespace_id; { let mut meta = write_txn.open_table(META_TABLE)?; let mut record = { @@ -339,6 +376,7 @@ impl MetadataStore { old_phase = record.phase; old_tags = record.tags.clone(); old_vector_slot = record.vector_slot; + old_namespace_id = record.namespace_id; // Strip content fields. record.summary = String::new(); @@ -367,8 +405,8 @@ impl MetadataStore { // Tombstoned records are not tracked in the bitmap since their // vector slots are freed. { - let mut pi = self.phase_index.write().unwrap(); - pi.remove(old_vector_slot, old_phase); + let mut pi = self.phase_index.write().unwrap_or_else(|e| e.into_inner()); + pi.remove(old_namespace_id, old_vector_slot, old_phase); } Ok(()) @@ -379,16 +417,18 @@ impl MetadataStore { /// Uses the in-memory roaring bitmap index. Requires a META_TABLE /// lookup per matching slot to resolve vector_slot -> MemoryId. pub fn ids_in_phase(&self, phase: DecayPhase) -> Result, StorageError> { - let slots: Vec = { - let pi = self.phase_index.read().unwrap(); - pi.bitmap(phase).iter().collect() + let ns_slots: Vec<(u32, u32)> = { + let pi = self.phase_index.read().unwrap_or_else(|e| e.into_inner()); + pi.all_slots_in_phase(phase) }; - if slots.is_empty() { + if ns_slots.is_empty() { return Ok(Vec::new()); } - let slot_set: std::collections::HashSet = slots.into_iter().collect(); + // Use (namespace_id, vector_slot) pairs for lookup to avoid + // cross-namespace collisions. + let slot_set: std::collections::HashSet<(u32, u32)> = ns_slots.into_iter().collect(); let mut ids = Vec::new(); let read_txn = self.db.begin_read()?; @@ -396,7 +436,7 @@ impl MetadataStore { for result in table.iter()? { let (key_bytes, value) = result?; let record = DiskRecord::from_bytes(value.value())?; - if slot_set.contains(&record.vector_slot) { + if slot_set.contains(&(record.namespace_id, record.vector_slot)) { let uuid = uuid::Uuid::from_slice(key_bytes.value()) .map_err(|_| StorageError::CorruptIndex("invalid UUID in meta table".into()))?; ids.push(MemoryId::from_uuid(uuid)); @@ -415,28 +455,30 @@ impl MetadataStore { &self, phase: DecayPhase, ) -> Result, StorageError> { - let slots: std::collections::HashSet = { - let pi = self.phase_index.read().unwrap(); - pi.bitmap(phase).iter().collect() + // Use (namespace_id, vector_slot) pairs for lookup to avoid + // cross-namespace collisions. + let ns_slots: std::collections::HashSet<(u32, u32)> = { + let pi = self.phase_index.read().unwrap_or_else(|e| e.into_inner()); + pi.all_slots_in_phase(phase).into_iter().collect() }; - if slots.is_empty() { + if ns_slots.is_empty() { return Ok(Vec::new()); } - let mut results = Vec::with_capacity(slots.len()); + let mut results = Vec::with_capacity(ns_slots.len()); let read_txn = self.db.begin_read()?; let table = read_txn.open_table(META_TABLE)?; for result in table.iter()? { let (key_bytes, value) = result?; let record = DiskRecord::from_bytes(value.value())?; - if slots.contains(&record.vector_slot) { + if ns_slots.contains(&(record.namespace_id, record.vector_slot)) { let uuid = uuid::Uuid::from_slice(key_bytes.value()) .map_err(|_| StorageError::CorruptIndex("invalid UUID in meta table".into()))?; results.push((MemoryId::from_uuid(uuid), record)); } - if results.len() == slots.len() { + if results.len() == ns_slots.len() { break; } } @@ -662,10 +704,15 @@ impl MetadataStore { // Phase 3: Update in-memory phase bitmap. { - let mut pi = self.phase_index.write().unwrap(); + let mut pi = self.phase_index.write().unwrap_or_else(|e| e.into_inner()); for (_, record, old_phase) in &updates { if *old_phase != record.phase { - pi.transition(record.vector_slot, *old_phase, record.phase); + pi.transition( + record.namespace_id, + record.vector_slot, + *old_phase, + record.phase, + ); } } @@ -700,7 +747,7 @@ impl MetadataStore { /// Called periodically (end of decay sweep, graceful shutdown). pub fn persist_phase_index(&self) -> Result<(), StorageError> { let bytes = { - let pi = self.phase_index.read().unwrap(); + let pi = self.phase_index.read().unwrap_or_else(|e| e.into_inner()); pi.to_bytes() }; let write_txn = self.db.begin_write()?; @@ -724,7 +771,7 @@ impl MetadataStore { for result in table.iter()? { let (_, value) = result?; let record = DiskRecord::from_bytes(value.value())?; - slot_phases.push((record.vector_slot, record.phase)); + slot_phases.push((record.namespace_id, record.vector_slot, record.phase)); } drop(table); drop(read_txn); @@ -741,7 +788,7 @@ impl MetadataStore { write_txn.commit()?; // Replace the in-memory copy. - *self.phase_index.write().unwrap() = rebuilt; + *self.phase_index.write().unwrap_or_else(|e| e.into_inner()) = rebuilt; Ok(()) } diff --git a/src/storage/text.rs b/src/storage/text.rs index 3a79b78..5d9b857 100644 --- a/src/storage/text.rs +++ b/src/storage/text.rs @@ -361,15 +361,21 @@ impl TextStore { &self.path } - /// Compact fulltext.dat by copying only live entries into a new file, - /// then atomically replacing the old file. + /// Phase 1 of fulltext.dat compaction: write a new file containing + /// only live entries, but do NOT rename or remove the marker. /// - /// `live_refs` contains (MemoryId, TextRef) for every DiskRecord - /// that still has a live text pointer. - pub fn compact( + /// The caller must update meta.db with the new offsets from the + /// returned `CompactionResult`, then call [`compact_finalize`] to + /// atomically swap the files. This two-phase approach ensures that + /// meta.db offsets are always consistent with the active fulltext.dat. + /// + /// Returns `(result, compacted)` where `compacted` is `true` if a + /// new file was written (fragmentation above threshold) and + /// `compact_finalize()` must be called afterwards. + pub fn compact_write_new_file( &mut self, live_refs: &[(MemoryId, TextRef)], - ) -> Result { + ) -> Result<(CompactionResult, bool), StorageError> { let parent_dir = self .path .parent() @@ -400,16 +406,19 @@ impl TextStore { threshold = format!("{:.0}%", COMPACTION_FRAGMENTATION_THRESHOLD * 100.0), "Skipping fulltext.dat compaction: fragmentation below threshold" ); - return Ok(CompactionResult { - old_size, - new_size: old_size, - entries_removed: 0, - entries_kept: live_refs.len(), - new_refs: live_refs - .iter() - .map(|&(id, text_ref)| (id, text_ref)) - .collect(), - }); + return Ok(( + CompactionResult { + old_size, + new_size: old_size, + entries_removed: 0, + entries_kept: live_refs.len(), + new_refs: live_refs + .iter() + .map(|&(id, text_ref)| (id, text_ref)) + .collect(), + }, + false, + )); } // Step 1: Write marker. @@ -462,20 +471,6 @@ impl TextStore { new_file.sync_all()?; drop(new_file); - // Step 5: Atomic rename. - fs::rename(&new_path, &self.path)?; - - // Step 6: Fsync directory. - fsync_dir(&parent_dir)?; - - // Step 7: Remove marker. - let _ = fs::remove_file(&marker_path); - fsync_dir(&parent_dir)?; - - // Step 8: Reopen file handle. - self.file = OpenOptions::new().read(true).write(true).open(&self.path)?; - self.write_pos = new_write_pos; - let entries_removed = total_possible.saturating_sub(live_refs.len()); let result = CompactionResult { @@ -486,15 +481,44 @@ impl TextStore { new_refs, }; - tracing::info!( - old_size = result.old_size, - new_size = result.new_size, - removed = result.entries_removed, - kept = result.entries_kept, - "Text.log compaction complete" - ); + Ok((result, true)) + } + + /// Phase 2 of fulltext.dat compaction: atomically rename the new + /// file over the old one, remove markers, and reopen the file handle. + /// + /// Must only be called after `compact_write_new_file` returned + /// `compacted = true` AND the caller has persisted updated offsets + /// to meta.db. + pub fn compact_finalize(&mut self) -> Result<(), StorageError> { + let parent_dir = self + .path + .parent() + .ok_or(StorageError::InvalidPath)? + .to_path_buf(); + let marker_path = parent_dir.join(".compacting"); + let committed_marker_path = parent_dir.join(".compaction-meta-committed"); + let new_path = parent_dir.join("fulltext.dat.new"); + + // Step 1: Atomic rename fulltext.dat.new -> fulltext.dat. + fs::rename(&new_path, &self.path)?; + + // Step 2: Fsync directory to persist the rename. + fsync_dir(&parent_dir)?; + + // Step 3: Remove markers. + let _ = fs::remove_file(&committed_marker_path); + let _ = fs::remove_file(&marker_path); + fsync_dir(&parent_dir)?; + + // Step 4: Reopen file handle with updated contents. + self.file = OpenOptions::new().read(true).write(true).open(&self.path)?; + // Determine new write position from file size. + self.write_pos = self.file.seek(SeekFrom::End(0))?; - Ok(result) + tracing::info!(new_size = self.write_pos, "Text.log compaction finalized"); + + Ok(()) } /// Approximate entry count by scanning forward from header. @@ -526,9 +550,23 @@ impl TextStore { /// Recover from an interrupted fulltext.dat compaction. /// /// Called once during startup validation, BEFORE opening the TextStore. +/// +/// Handles two compaction states using marker files: +/// +/// - `.compacting` alone (no `.compaction-meta-committed`): +/// meta.db has NOT been updated with new offsets yet. Delete the +/// incomplete `fulltext.dat.new` and revert to the original file. +/// +/// - `.compacting` AND `.compaction-meta-committed`: +/// meta.db HAS been updated with new offsets. Complete the rename +/// of `fulltext.dat.new` -> `fulltext.dat` so the file matches +/// the committed offsets. If the rename already happened (no .new +/// file), just clean up the markers. pub fn recover_text_compaction(db_dir: &Path) -> Result<(), StorageError> { let marker_path = db_dir.join(".compacting"); + let committed_marker_path = db_dir.join(".compaction-meta-committed"); let new_file_path = db_dir.join("fulltext.dat.new"); + let fulltext_path = db_dir.join("fulltext.dat"); if !marker_path.exists() { // No compaction was in progress. @@ -540,15 +578,33 @@ pub fn recover_text_compaction(db_dir: &Path) -> Result<(), StorageError> { marker_path.display() ); - // If fulltext.dat.new exists, it is incomplete or un-swapped. - // meta.db still has old offsets. Safe to delete .new and revert. - if new_file_path.exists() { - tracing::warn!("Deleting incomplete fulltext.dat.new"); - fs::remove_file(&new_file_path)?; + if committed_marker_path.exists() { + // meta.db has already been updated with new offsets. + // We must complete the rename so fulltext.dat matches. + if new_file_path.exists() { + tracing::warn!( + "meta.db offsets committed; completing rename of fulltext.dat.new -> fulltext.dat" + ); + fs::rename(&new_file_path, &fulltext_path)?; + fsync_dir(db_dir)?; + } else { + // Rename already happened before crash; just clean up markers. + tracing::info!( + "meta.db offsets committed and rename already complete; cleaning up markers" + ); + } + } else { + // meta.db has NOT been updated. Safe to delete .new and revert + // to the original fulltext.dat with its original offsets. + if new_file_path.exists() { + tracing::warn!("Deleting incomplete fulltext.dat.new"); + fs::remove_file(&new_file_path)?; + } } - // Remove the marker. - fs::remove_file(&marker_path)?; + // Remove all markers. + let _ = fs::remove_file(&committed_marker_path); + let _ = fs::remove_file(&marker_path); fsync_dir(db_dir)?; tracing::info!("Compaction recovery complete");