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

engram_core/
server.rs
