engram_core/
server.rs

1use std::path::{Path, PathBuf};
2use std::sync::{Arc, Mutex, RwLock};
3
4use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
5use tokio::net::UnixListener;
6use tokio::signal::unix::{SignalKind, signal};
7
8use engram_embeddings::Embedder;
9use engram_llm_client::{EmbeddingProvider, TextGenerator};
10use engram_router::Router;
11use engram_storage::Database;
12
13use crate::config::{self, Config};
14use crate::dispatch;
15use crate::error::CoreError;
16use crate::indexes::IndexSet;
17use crate::lock_helpers;
18use crate::protocol::{JsonRequest, JsonResponse};
19
20const PROJECT_DB_RELATIVE: &str = ".engram/engram.db";
21const PROJECT_SOCKET_RELATIVE: &str = ".engram/engram.sock";
22
23const ROUTER_DEFAULT_ALPHA: f32 = 0.1;
24const ROUTER_DEFAULT_EPSILON: f32 = 0.15;
25
26pub struct ServerState {
27    pub database: Mutex<Database>,
28    pub indexes: RwLock<IndexSet>,
29    pub embedder: Embedder,
30    pub router: Mutex<Router>,
31    pub config: Config,
32    pub database_path: String,
33    pub embedding_provider: Arc<dyn EmbeddingProvider + Send + Sync>,
34    pub text_generator: Option<Arc<dyn TextGenerator + Send + Sync>>,
35}
36
37pub async fn run(config: Config) -> Result<(), CoreError> {
38    let cwd = std::env::current_dir()
39        .map_err(|error| CoreError::SocketError(format!("cwd unavailable: {error}")))?;
40    let home = home_dir_or_error()?;
41    let project_dir = config::resolve_project_dir(&cwd, None)?;
42    // See execute_command in main.rs for why this construction has to run on
43    // the blocking pool. reqwest::blocking::ClientBuilder::build internally
44    // spins up and drops a current-thread tokio runtime; doing that on the
45    // outer multi-threaded runtime's worker panics.
46    let state = {
47        let config = config.clone();
48        let project_dir = project_dir.clone();
49        let home = home.clone();
50        tokio::task::spawn_blocking(move || initialize_state(&config, &project_dir, &home))
51            .await
52            .map_err(|error| CoreError::SocketError(error.to_string()))??
53    };
54    {
55        let database = lock_helpers::lock_db(&state);
56        crate::migrations::embedding_model_v1::check(
57            &database,
58            state.embedding_provider.model_name(),
59        )?;
60    }
61    warn_missing_api_keys(&config);
62    let shared_state = Arc::new(state);
63    let socket_path = resolve_socket_path(&project_dir, &config)?;
64    cleanup_stale_socket(&socket_path);
65    let listener = bind_listener(&socket_path)?;
66    accept_loop(listener, shared_state).await
67}
68
69pub fn check_legacy_database(project_dir: &Path, home_dir: &Path) -> Result<(), CoreError> {
70    let project_db = project_dir.join(PROJECT_DB_RELATIVE);
71    let home_db = home_dir.join(PROJECT_DB_RELATIVE);
72    if !project_db.exists() && home_db.exists() {
73        return Err(CoreError::LegacyDatabaseDetected {
74            legacy_path: home_db.to_string_lossy().to_string(),
75            project_path: project_dir.to_string_lossy().to_string(),
76        });
77    }
78    Ok(())
79}
80
81pub(crate) fn initialize_state(
82    config: &Config,
83    project_dir: &Path,
84    home_dir: &Path,
85) -> Result<ServerState, CoreError> {
86    crate::config::validate_dedup_threshold(config.deduplication.threshold)?;
87    crate::config::validate_search_config(&config.search)?;
88    crate::config::validate_fts_similarity_floor(config.consolidation.fts_similarity_floor)?;
89    check_legacy_database(project_dir, home_dir)?;
90    let database_path = resolve_database_path(project_dir);
91    let database = Database::open(&database_path)?;
92    crate::migrations::run_pending(&database)?;
93    let embedding_provider: Arc<dyn EmbeddingProvider + Send + Sync> =
94        Arc::from(config.build_embedding_provider()?);
95    let text_generator: Option<Arc<dyn TextGenerator + Send + Sync>> =
96        match config.build_text_generator() {
97            Ok(generator) => Some(Arc::from(generator)),
98            Err(error) => {
99                eprintln!("warning: text generator unavailable, running in degraded mode: {error}");
100                None
101            }
102        };
103    let hnsw_config = config.clone();
104    let indexes = crate::persistence::load_or_rebuild(
105        &resolve_index_directory(&database_path),
106        &database,
107        || hnsw_config.build_hnsw_params(),
108    )?;
109    let embedder = Embedder::new(config.embedding.hyde_threshold);
110    let router = Router::new(ROUTER_DEFAULT_ALPHA, ROUTER_DEFAULT_EPSILON);
111    Ok(ServerState {
112        database: Mutex::new(database),
113        indexes: RwLock::new(indexes),
114        embedder,
115        router: Mutex::new(router),
116        config: config.clone(),
117        database_path,
118        embedding_provider,
119        text_generator,
120    })
121}
122
123fn resolve_database_path(project_dir: &Path) -> String {
124    if let Ok(value) = std::env::var("ENGRAM_DB_PATH") {
125        return value;
126    }
127    project_dir
128        .join(PROJECT_DB_RELATIVE)
129        .to_string_lossy()
130        .to_string()
131}
132
133fn resolve_socket_path(project_dir: &Path, config: &Config) -> Result<String, CoreError> {
134    if let Ok(value) = std::env::var("ENGRAM_SOCKET_PATH") {
135        return Ok(value);
136    }
137    // Prefer project socket when the `.engram/` dir exists (walk-up already validated this).
138    if project_dir.join(".engram").is_dir() {
139        return Ok(project_dir
140            .join(PROJECT_SOCKET_RELATIVE)
141            .to_string_lossy()
142            .to_string());
143    }
144    match config.server.socket_path.as_deref() {
145        Some(fallback) => Ok(expand_tilde(fallback)),
146        None => Err(CoreError::SocketError(
147            "no socket path available: no project .engram/ marker found, \
148             ENGRAM_SOCKET_PATH not set, and server.socket_path not configured"
149                .into(),
150        )),
151    }
152}
153
154fn home_dir_or_error() -> Result<PathBuf, CoreError> {
155    match std::env::var("HOME") {
156        Ok(value) if !value.is_empty() => Ok(PathBuf::from(value)),
157        _ => Err(CoreError::InitFailed(
158            "HOME environment variable not set".into(),
159        )),
160    }
161}
162
163pub(crate) fn resolve_index_directory(database_path: &str) -> String {
164    std::path::Path::new(database_path)
165        .parent()
166        .map(|parent| parent.to_string_lossy().to_string())
167        .unwrap_or_else(|| ".".to_string())
168}
169
170fn cleanup_stale_socket(socket_path: &str) {
171    let _ = std::fs::remove_file(socket_path);
172}
173
174fn bind_listener(socket_path: &str) -> Result<UnixListener, CoreError> {
175    if let Some(parent) = std::path::Path::new(socket_path).parent() {
176        std::fs::create_dir_all(parent)
177            .map_err(|error| CoreError::SocketError(error.to_string()))?;
178    }
179    UnixListener::bind(socket_path).map_err(|error| CoreError::SocketError(error.to_string()))
180}
181
182async fn accept_loop(listener: UnixListener, state: Arc<ServerState>) -> Result<(), CoreError> {
183    let mut sigterm =
184        signal(SignalKind::terminate()).map_err(|e| CoreError::SocketError(e.to_string()))?;
185    spawn_background_reindex(Arc::clone(&state));
186    loop {
187        tokio::select! {
188            accept_result = listener.accept() => {
189                let (stream, _address) = accept_result
190                    .map_err(|error| CoreError::SocketError(error.to_string()))?;
191                let client_state = Arc::clone(&state);
192                tokio::spawn(async move {
193                    if let Err(error) = handle_client(stream, client_state).await {
194                        eprintln!("client error: {error}");
195                    }
196                });
197            }
198            _ = sigterm.recv() => {
199                return save_indexes_on_shutdown(&state);
200            }
201        }
202    }
203}
204
205fn spawn_background_reindex(state: Arc<ServerState>) {
206    let interval_secs = state.config.server.reindex_interval_secs;
207    if interval_secs == 0 {
208        return;
209    }
210    tokio::spawn(async move {
211        let mut interval = tokio::time::interval(std::time::Duration::from_secs(interval_secs));
212        interval.tick().await; // skip first immediate tick
213        loop {
214            interval.tick().await;
215            let reindex_state = Arc::clone(&state);
216            let _ = tokio::task::spawn_blocking(move || {
217                reindex_unindexed_memories(&reindex_state);
218            })
219            .await;
220        }
221    });
222}
223
224/// Reindex unindexed memories from SQLite into the HNSW index.
225///
226/// **Unstable:** exposed for latency benchmarking only
227/// (`crates/engram-core/tests/search_latency_bench.rs`). Not part of the public
228/// API contract; may change or be removed without notice.
229#[doc(hidden)]
230pub fn reindex_unindexed_memories(state: &ServerState) {
231    let database = lock_helpers::lock_db(state);
232    let unindexed = match query_unindexed_ids(&database) {
233        Ok(ids) => ids,
234        Err(_) => return,
235    };
236    if unindexed.is_empty() {
237        return;
238    }
239    let mut indexes = lock_helpers::write_indexes(state);
240    for memory_id in &unindexed {
241        let _ = reindex_single_memory(&database, &mut indexes, memory_id);
242    }
243}
244
245fn query_unindexed_ids(database: &Database) -> Result<Vec<String>, CoreError> {
246    let mut statement = database
247        .connection()
248        .prepare(
249            "SELECT id FROM memories WHERE indexed = 0 \
250             AND embedding_context IS NOT NULL LIMIT 100",
251        )
252        .map_err(|e| CoreError::Storage(engram_storage::StorageError::from(e)))?;
253    let rows = statement
254        .query_map([], |row| row.get::<_, String>(0))
255        .map_err(|e| CoreError::Storage(engram_storage::StorageError::from(e)))?;
256    let ids: Vec<String> = rows.flatten().collect();
257    Ok(ids)
258}
259
260fn reindex_single_memory(
261    database: &Database,
262    indexes: &mut IndexSet,
263    memory_id: &str,
264) -> Result<(), CoreError> {
265    let memory = database.get_memory(memory_id)?;
266    let embedding = crate::persistence::extract_embeddings_from_memory(&memory)?;
267    let id_hash = crate::persistence::hash_string_to_u64(memory_id);
268    let rng_value = crate::persistence::deterministic_rng(id_hash);
269    if !indexes.contains(id_hash) {
270        indexes.insert_atomic(id_hash, memory_id, &embedding, rng_value)?;
271    }
272    database.set_memory_indexed(memory_id, true)?;
273    Ok(())
274}
275
276async fn handle_client(
277    stream: tokio::net::UnixStream,
278    state: Arc<ServerState>,
279) -> Result<(), CoreError> {
280    let (reader, mut writer) = stream.into_split();
281    let mut lines = BufReader::new(reader).lines();
282    while let Some(line) = lines
283        .next_line()
284        .await
285        .map_err(|error| CoreError::SocketError(error.to_string()))?
286    {
287        let response = process_request_line(&line, &state).await;
288        let serialized = serde_json::to_string(&response)
289            .map_err(|error| CoreError::SocketError(error.to_string()))?;
290        writer
291            .write_all(serialized.as_bytes())
292            .await
293            .map_err(|error| CoreError::SocketError(error.to_string()))?;
294        writer
295            .write_all(b"\n")
296            .await
297            .map_err(|error| CoreError::SocketError(error.to_string()))?;
298    }
299    Ok(())
300}
301
302async fn process_request_line(line: &str, state: &Arc<ServerState>) -> JsonResponse {
303    let request: JsonRequest = match serde_json::from_str(line) {
304        Ok(parsed) => parsed,
305        Err(error) => {
306            return JsonResponse::error(String::new(), 4000, format!("invalid json: {error}"));
307        }
308    };
309    let request_id = request.id.clone();
310    match dispatch::route(&request.method, state, request.params).await {
311        Ok(data) => JsonResponse::success(request_id, data),
312        Err(error) => JsonResponse::error(request_id, error_code(&error), error.to_string()),
313    }
314}
315
316fn error_code(error: &CoreError) -> u32 {
317    match error {
318        CoreError::ConfigNotFound => 6001,
319        CoreError::ConfigParseError(_) => 6002,
320        CoreError::InvalidProvider(_) => 6003,
321        CoreError::IndexCorrupted(_) => 6004,
322        CoreError::RebuildFailed(_) => 6005,
323        CoreError::SocketError(_) => 6006,
324        CoreError::DispatchError(_) => 6007,
325        CoreError::ConfigReadOnly => 6008,
326        CoreError::ExportFailed(_) => 6009,
327        CoreError::ImportVersionMismatch(_) => 6010,
328        CoreError::ImportFailed(_) => 6011,
329        CoreError::InitFailed(_) => 6012,
330        CoreError::TrainerFailed(_) => 6013,
331        CoreError::TrainerTimeout => 6014,
332        CoreError::TrainerMalformedOutput(_) => 6015,
333        CoreError::ProjectDirNotFound => 6016,
334        CoreError::LegacyDatabaseDetected { .. } => 6017,
335        CoreError::MigrationSourceNotFound => 6018,
336        CoreError::MigrationFailed(_) => 6019,
337        CoreError::EmbeddingModelMismatch { .. } => 6020,
338        CoreError::IndexHashCollision { .. } => 6021,
339        CoreError::ConfigValidation(_) => 6022,
340        CoreError::Storage(_) => 1000,
341        CoreError::Hnsw(_) => 3000,
342        CoreError::Api(_) => 2000,
343        CoreError::Consolidation(_) => 5000,
344    }
345}
346
347fn save_indexes_on_shutdown(state: &Arc<ServerState>) -> Result<(), CoreError> {
348    let index_directory = resolve_index_directory(&state.database_path);
349    let indexes = lock_helpers::read_indexes(state);
350    crate::persistence::save_to_disk(&index_directory, &indexes)
351}
352
353fn warn_missing_api_keys(config: &Config) {
354    let has_embedding_key = config
355        .embedding
356        .api_key
357        .as_deref()
358        .is_some_and(|k| !k.is_empty());
359    let has_llm_key = config.llm.api_key.as_deref().is_some_and(|k| !k.is_empty());
360
361    if !has_embedding_key
362        && config.embedding.provider != "deterministic"
363        && config.embedding.provider != "ollama"
364    {
365        eprintln!(
366            "warning: {} embedding provider configured without api key (set ENGRAM_VOYAGE_API_KEY)",
367            config.embedding.provider
368        );
369    }
370    if !has_llm_key && config.llm.provider != "local" && config.llm.provider != "ollama" {
371        eprintln!(
372            "warning: {} llm provider configured without api key (set ENGRAM_OPENAI_API_KEY)",
373            config.llm.provider
374        );
375    }
376}
377
378fn expand_tilde(path: &str) -> String {
379    if let Some(rest) = path.strip_prefix("~/")
380        && let Ok(home) = std::env::var("HOME")
381    {
382        return format!("{home}/{rest}");
383    }
384    path.to_string()
385}
engram_core/server.rs

engram_core/
server.rs
