Skip to main content

engram_core/
store_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_storage::Memory;
7
8use crate::error::CoreError;
9use crate::lock_helpers;
10use crate::persistence::f32_vec_to_bytes;
11use crate::server::ServerState;
12use crate::tags_normalize::{TagsInput, normalize_tags};
13use crate::timestamp::current_utc_timestamp;
14
15const MAX_FIELD_LENGTH: usize = 10_000;
16
17#[derive(Deserialize)]
18struct StoreParams {
19    memory_type: String,
20    context: String,
21    action: String,
22    result: String,
23    tags: Option<TagsInput>,
24    project: Option<String>,
25}
26
27pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
28    let mut parsed: StoreParams = serde_json::from_value(params)
29        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
30    validate_field_length("context", &parsed.context)?;
31    validate_field_length("action", &parsed.action)?;
32    validate_field_length("result", &parsed.result)?;
33    let normalized_tags = normalize_tags(parsed.tags.take());
34    let memory_id = uuid::Uuid::new_v4().to_string();
35    let timestamp = current_utc_timestamp();
36    let embedding = compute_embedding(state, &parsed).await?;
37    if let Some(existing_id) = find_duplicate_memory(state, &embedding).await? {
38        return deduplicate_into(state, &existing_id, &timestamp).await;
39    }
40    let memory = build_memory(&memory_id, &timestamp, &parsed, normalized_tags, &embedding);
41    let indexed = persist_memory(state, &memory, &memory_id, &embedding).await?;
42    Ok(json!({ "id": memory_id, "indexed": indexed, "deduplicated": false }))
43}
44
45/// Searches the HNSW index for an existing near-duplicate of `embedding` under
46/// the all-three-fields policy. The new memory is not yet in the index, so this
47/// never self-matches. A search error degrades gracefully: it logs one warning
48/// and returns `Ok(None)`, letting the insert proceed rather than failing the
49/// store over a non-fatal dedup hiccup.
50async fn find_duplicate_memory(
51    state: &Arc<ServerState>,
52    embedding: &engram_embeddings::ThreeFieldEmbedding,
53) -> Result<Option<String>, CoreError> {
54    let threshold = state.config.deduplication.threshold;
55    let embedding_owned = engram_embeddings::ThreeFieldEmbedding {
56        context: embedding.context.clone(),
57        action: embedding.action.clone(),
58        result: embedding.result.clone(),
59    };
60    let state_clone = Arc::clone(state);
61    tokio::task::spawn_blocking(move || {
62        let indexes = lock_helpers::read_indexes(&state_clone);
63        match indexes.find_duplicate(&embedding_owned, threshold) {
64            Ok(Some((node, _similarity))) => Ok(indexes.resolve_node_id(node).map(str::to_string)),
65            Ok(None) => Ok(None),
66            Err(hnsw) => {
67                eprintln!("warning: dedup search failed, inserting without dedup: {hnsw}");
68                Ok(None)
69            }
70        }
71    })
72    .await
73    .map_err(|error| CoreError::SocketError(error.to_string()))?
74}
75
76/// Records a hit against the existing duplicate (bumping its `used_count` and
77/// `last_used_at`) instead of inserting a new row, then reports the merge.
78async fn deduplicate_into(
79    state: &Arc<ServerState>,
80    existing_id: &str,
81    timestamp: &str,
82) -> Result<Value, CoreError> {
83    let existing_id_owned = existing_id.to_string();
84    let timestamp_owned = timestamp.to_string();
85    let state_clone = Arc::clone(state);
86    tokio::task::spawn_blocking(move || {
87        lock_helpers::lock_db(&state_clone).touch_memory(&existing_id_owned, &timestamp_owned)?;
88        Ok::<(), CoreError>(())
89    })
90    .await
91    .map_err(|error| CoreError::SocketError(error.to_string()))??;
92    Ok(json!({
93        "id": existing_id,
94        "indexed": true,
95        "deduplicated": true,
96        "merged_into": existing_id,
97    }))
98}
99
100async fn compute_embedding(
101    state: &Arc<ServerState>,
102    params: &StoreParams,
103) -> Result<engram_embeddings::ThreeFieldEmbedding, CoreError> {
104    let context = params.context.clone();
105    let action = params.action.clone();
106    let result = params.result.clone();
107    let state_clone = Arc::clone(state);
108    tokio::task::spawn_blocking(move || {
109        let provider = state_clone.embedding_provider.as_ref();
110        let text_gen_ref = state_clone
111            .text_generator
112            .as_deref()
113            .map(|generator| generator as &dyn engram_llm_client::TextGenerator);
114        state_clone
115            .embedder
116            .embed_fields(&context, &action, &result, provider, text_gen_ref)
117            .map_err(|error| {
118                CoreError::Api(engram_llm_client::ApiError::EmbeddingApiUnavailable(
119                    error.to_string(),
120                ))
121            })
122    })
123    .await
124    .map_err(|error| CoreError::SocketError(error.to_string()))?
125}
126
127fn build_memory(
128    memory_id: &str,
129    timestamp: &str,
130    params: &StoreParams,
131    normalized_tags: Option<String>,
132    embedding: &engram_embeddings::ThreeFieldEmbedding,
133) -> Memory {
134    Memory {
135        id: memory_id.to_string(),
136        memory_type: params.memory_type.clone(),
137        context: params.context.clone(),
138        action: params.action.clone(),
139        result: params.result.clone(),
140        score: 0.0,
141        embedding_context: Some(f32_vec_to_bytes(&embedding.context)),
142        embedding_action: Some(f32_vec_to_bytes(&embedding.action)),
143        embedding_result: Some(f32_vec_to_bytes(&embedding.result)),
144        indexed: false,
145        tags: normalized_tags,
146        project: params.project.clone(),
147        parent_id: None,
148        source_ids: None,
149        insight_type: None,
150        created_at: timestamp.to_string(),
151        updated_at: timestamp.to_string(),
152        used_count: 0,
153        last_used_at: None,
154        superseded_by: None,
155    }
156}
157
158/// Persists the memory to SQLite (the source of truth, written with
159/// `indexed=false`) and then attempts the HNSW index write. Returns the
160/// truthful final `indexed` state: `true` only after the HNSW write is
161/// confirmed and the row is marked indexed. A transient HNSW failure leaves
162/// the row at `indexed=false` for background reindex to recover; a hash
163/// collision is non-recoverable and propagates.
164async fn persist_memory(
165    state: &Arc<ServerState>,
166    memory: &Memory,
167    memory_id: &str,
168    embedding: &engram_embeddings::ThreeFieldEmbedding,
169) -> Result<bool, CoreError> {
170    insert_row(state, memory).await?;
171    match attempt_index(state, memory_id, embedding).await? {
172        Ok(()) => {
173            mark_indexed(state, memory_id).await?;
174            Ok(true)
175        }
176        Err(collision @ CoreError::IndexHashCollision { .. }) => Err(collision),
177        Err(transient) => {
178            eprintln!(
179                "warning: HNSW index write failed for {memory_id}, \
180                 left unindexed for background reindex: {transient}"
181            );
182            Ok(false)
183        }
184    }
185}
186
187async fn insert_row(state: &Arc<ServerState>, memory: &Memory) -> Result<(), CoreError> {
188    let memory_owned = memory.clone();
189    let state_db = Arc::clone(state);
190    tokio::task::spawn_blocking(move || {
191        let database = lock_helpers::lock_db(&state_db);
192        database.insert_memory(&memory_owned)?;
193        Ok::<(), CoreError>(())
194    })
195    .await
196    .map_err(|error| CoreError::SocketError(error.to_string()))?
197}
198
199async fn attempt_index(
200    state: &Arc<ServerState>,
201    memory_id: &str,
202    embedding: &engram_embeddings::ThreeFieldEmbedding,
203) -> Result<Result<(), CoreError>, CoreError> {
204    let hashed_id = hash_string_to_u64(memory_id);
205    let rng_value = deterministic_rng(hashed_id);
206    let memory_id_owned = memory_id.to_string();
207    let embedding_owned = engram_embeddings::ThreeFieldEmbedding {
208        context: embedding.context.clone(),
209        action: embedding.action.clone(),
210        result: embedding.result.clone(),
211    };
212    let state_idx = Arc::clone(state);
213    tokio::task::spawn_blocking(move || {
214        let mut indexes = lock_helpers::write_indexes(&state_idx);
215        indexes.insert_atomic(hashed_id, &memory_id_owned, &embedding_owned, rng_value)
216    })
217    .await
218    .map_err(|error| CoreError::SocketError(error.to_string()))
219}
220
221async fn mark_indexed(state: &Arc<ServerState>, memory_id: &str) -> Result<(), CoreError> {
222    let memory_id_owned = memory_id.to_string();
223    let state_db = Arc::clone(state);
224    tokio::task::spawn_blocking(move || {
225        let database = lock_helpers::lock_db(&state_db);
226        database.set_memory_indexed(&memory_id_owned, true)?;
227        Ok::<(), CoreError>(())
228    })
229    .await
230    .map_err(|error| CoreError::SocketError(error.to_string()))?
231}
232
233const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
234const FNV_PRIME: u64 = 0x0100_0000_01b3;
235const DETERMINISTIC_RNG_MULTIPLIER: u64 = 0x9e37_79b9_7f4a_7c15;
236
237fn hash_string_to_u64(value: &str) -> u64 {
238    let mut hash: u64 = FNV_OFFSET_BASIS;
239    for byte in value.as_bytes() {
240        hash ^= *byte as u64;
241        hash = hash.wrapping_mul(FNV_PRIME);
242    }
243    hash
244}
245
246fn deterministic_rng(id: u64) -> f64 {
247    let mixed = id.wrapping_mul(DETERMINISTIC_RNG_MULTIPLIER);
248    (mixed >> 11) as f64 / (1u64 << 53) as f64
249}
250
251fn validate_field_length(field_name: &str, value: &str) -> Result<(), CoreError> {
252    if value.len() > MAX_FIELD_LENGTH {
253        return Err(CoreError::DispatchError(format!(
254            "{field_name} exceeds maximum length of {MAX_FIELD_LENGTH} bytes"
255        )));
256    }
257    Ok(())
258}