Skip to main content

mentedb_server/
handlers_ops.rs

1//! Operations endpoints: maintenance, enrichment, documents, backups, database management, readiness.
2//!
3//! Per database (mounted under `/v1` and `/v1/db/{db}`):
4//!
5//! - `POST /maintenance/decay`, `/maintenance/archive`, `/maintenance/consolidate`,
6//!   `/maintenance/sweep`: the rule based upkeep jobs, on demand.
7//! - `POST /enrich`: the LLM enrichment pipeline (extraction, entity linking,
8//!   communities, user model).
9//! - `POST /documents`, `DELETE /documents/{document_id}`: load a long text as
10//!   overlapping chunks, and forget it again.
11//! - `POST /backup`, `GET /backups` (admin): snapshot the database into a
12//!   single backup file and list those snapshots.
13//!
14//! Global:
15//!
16//! - `GET`/`POST /v1/databases`, `DELETE /v1/databases/{name}` (admin).
17//! - `GET /v1/ready`: readiness, separate from `/v1/health` liveness.
18//!
19//! There is deliberately no restore endpoint: replacing the files of a live,
20//! open database is unsafe. Restore is offline: stop the server, then call
21//! `mentedb::storage::backup::BackupManager::restore_backup(file, dir)` with
22//! an empty target directory (the data directory root for the default
23//! database, `<data>/dbs/<name>` for a named one) and start the server again.
24
25use std::collections::HashMap;
26use std::path::{Path as FsPath, PathBuf};
27use std::sync::{Arc, Mutex, OnceLock, Weak};
28use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
29
30use axum::extract::{FromRequestParts, Path, State};
31use axum::http::StatusCode;
32use axum::http::request::Parts;
33use axum::response::{IntoResponse, Response};
34use axum::routing::{delete, get, post};
35use axum::{Extension, Json, Router};
36use mentedb::consolidation::archival::ArchivalDecision;
37use mentedb::storage::backup::BackupManager;
38use mentedb_core::MemoryNode;
39use mentedb_core::memory::{AttributeValue, MemoryType};
40use mentedb_core::types::{AgentId, MemoryId, SpaceId, UserId};
41use mentedb_embedding::provider::EmbeddingProvider;
42use serde_json::{Value, json};
43use uuid::Uuid;
44
45use crate::auth::AuthenticatedAgent;
46use crate::databases::{DEFAULT_DB, Db, validate_name};
47use crate::error::ApiError;
48use crate::state::AppState;
49
50// ---------------------------------------------------------------------------
51// Configuration
52// ---------------------------------------------------------------------------
53
54/// How documents are split into chunks.
55#[derive(Clone, Debug)]
56pub struct DocumentConfig {
57    /// Target chunk length in characters.
58    pub chunk_chars: usize,
59    /// Characters carried from the end of one chunk to the start of the next.
60    pub overlap_chars: usize,
61    /// Smallest `chunk_chars` a request may ask for.
62    pub min_chunk_chars: usize,
63    /// Largest `chunk_chars` a request may ask for.
64    pub max_chunk_chars: usize,
65    /// Largest overlap as a fraction of the chunk length.
66    pub max_overlap_ratio: f32,
67    /// Most chunks one document may produce.
68    pub max_chunks: usize,
69    /// Characters that end a sentence when a paragraph must be split.
70    pub sentence_terminators: Vec<char>,
71}
72
73impl Default for DocumentConfig {
74    fn default() -> Self {
75        Self {
76            chunk_chars: 1500,
77            overlap_chars: 200,
78            min_chunk_chars: 100,
79            max_chunk_chars: 8000,
80            max_overlap_ratio: 0.5,
81            max_chunks: 10_000,
82            sentence_terminators: vec!['.', '!', '?', '\u{3002}', '\u{ff01}', '\u{ff1f}'],
83        }
84    }
85}
86
87/// How `/v1/ready` probes its dependencies.
88#[derive(Clone, Debug)]
89pub struct ReadinessConfig {
90    /// How long a successful or failed embed check is reused before the
91    /// embedder is asked again, so frequent probes stay cheap.
92    pub embed_check_interval: Duration,
93    /// Text embedded by the probe.
94    pub probe_text: String,
95    /// Longest any single check may take before it counts as failed.
96    pub check_timeout: Duration,
97}
98
99impl Default for ReadinessConfig {
100    fn default() -> Self {
101        Self {
102            embed_check_interval: Duration::from_secs(30),
103            probe_text: "readiness probe".to_string(),
104            check_timeout: Duration::from_secs(5),
105        }
106    }
107}
108
109/// Settings for the operations endpoints.
110#[derive(Clone, Debug)]
111pub struct OpsConfig {
112    /// Where backups are written (`MENTEDB_BACKUP_DIR`). `None` means
113    /// `<data_dir>/backups`. Each database gets its own subdirectory.
114    pub backup_dir: Option<PathBuf>,
115    /// Smallest cluster of near duplicates worth consolidating.
116    pub consolidation_min_cluster_size: usize,
117    /// Cosine similarity at which memories are grouped for consolidation.
118    pub consolidation_similarity: f32,
119    /// Memory ids handed to one paraphrase dedup pass, so a large database
120    /// is swept in bounded steps.
121    pub dedup_batch_size: usize,
122    /// Most memories one archive run may forget.
123    pub max_forgets_per_archive: usize,
124    pub documents: DocumentConfig,
125    pub readiness: ReadinessConfig,
126}
127
128impl Default for OpsConfig {
129    fn default() -> Self {
130        Self {
131            backup_dir: None,
132            consolidation_min_cluster_size: crate::maintenance::MIN_CLUSTER_SIZE,
133            consolidation_similarity: crate::maintenance::SIMILARITY_THRESHOLD,
134            dedup_batch_size: 1000,
135            max_forgets_per_archive: crate::maintenance::MAX_FORGETS_PER_SWEEP,
136            documents: DocumentConfig::default(),
137            readiness: ReadinessConfig::default(),
138        }
139    }
140}
141
142impl OpsConfig {
143    /// Defaults plus `MENTEDB_BACKUP_DIR`.
144    pub fn from_env() -> Self {
145        Self {
146            backup_dir: std::env::var_os("MENTEDB_BACKUP_DIR")
147                .filter(|v| !v.is_empty())
148                .map(PathBuf::from),
149            ..Self::default()
150        }
151    }
152}
153
154/// The operations settings, read from the environment once.
155pub fn ops_config() -> &'static OpsConfig {
156    static CONFIG: OnceLock<OpsConfig> = OnceLock::new();
157    CONFIG.get_or_init(OpsConfig::from_env)
158}
159
160// ---------------------------------------------------------------------------
161// Routes
162// ---------------------------------------------------------------------------
163
164/// Routes on one database, mounted under `/v1` and `/v1/db/{db}`.
165pub fn routes() -> Router<Arc<AppState>> {
166    Router::new()
167        .route("/maintenance/decay", post(maintenance_decay))
168        .route("/maintenance/archive", post(maintenance_archive))
169        .route("/maintenance/consolidate", post(maintenance_consolidate))
170        .route("/maintenance/sweep", post(maintenance_sweep))
171        .route("/enrich", post(enrich))
172        .route("/documents", post(store_document))
173        .route("/documents/{document_id}", delete(delete_document))
174        .route("/backup", post(create_backup))
175        .route("/backups", get(list_backups))
176}
177
178/// Routes outside any one database (database management, readiness).
179pub fn global_routes() -> Router<Arc<AppState>> {
180    Router::new()
181        .route("/v1/databases", get(list_databases).post(create_database))
182        .route("/v1/databases/{name}", delete(delete_database))
183        .route("/v1/ready", get(ready))
184}
185
186// ---------------------------------------------------------------------------
187// Helpers
188// ---------------------------------------------------------------------------
189
190/// Admits only admins: a token with the admin claim, or the configured admin
191/// key in `x-api-key` / `Authorization: Bearer`. Listed before [`Db`] in a
192/// handler so a refused request never opens or creates a database.
193pub struct AdminGuard;
194
195impl FromRequestParts<Arc<AppState>> for AdminGuard {
196    type Rejection = ApiError;
197
198    async fn from_request_parts(
199        parts: &mut Parts,
200        state: &Arc<AppState>,
201    ) -> Result<Self, ApiError> {
202        if parts
203            .extensions
204            .get::<AuthenticatedAgent>()
205            .is_some_and(|a| a.admin)
206        {
207            return Ok(Self);
208        }
209        crate::auth::admin_authorized(&parts.headers, &state.admin_key)?;
210        Ok(Self)
211    }
212}
213
214/// Run blocking engine work off the async workers.
215async fn blocking<T, F>(what: &'static str, f: F) -> Result<T, ApiError>
216where
217    T: Send + 'static,
218    F: FnOnce() -> Result<T, ApiError> + Send + 'static,
219{
220    tokio::task::spawn_blocking(f)
221        .await
222        .map_err(|e| ApiError::Internal(format!("{what} task failed: {e}")))?
223}
224
225fn internal(what: &str) -> impl FnOnce(mentedb_core::MenteError) -> ApiError + '_ {
226    move |e| ApiError::Internal(format!("{what} failed: {e}"))
227}
228
229fn now_micros() -> u64 {
230    SystemTime::now()
231        .duration_since(UNIX_EPOCH)
232        .unwrap_or_default()
233        .as_micros() as u64
234}
235
236// ---------------------------------------------------------------------------
237// POST /maintenance/*
238// ---------------------------------------------------------------------------
239
240/// Age salience. The engine derives decay on read, so this persists nothing
241/// and reports what the engine reports.
242async fn maintenance_decay(db: Db) -> Result<Json<Value>, ApiError> {
243    let handle = db.handle.clone();
244    let decayed = blocking("decay", move || {
245        handle.apply_decay_global().map_err(internal("decay"))
246    })
247    .await?;
248    Ok(Json(json!({ "database": db.name, "decayed": decayed })))
249}
250
251/// Evaluate every memory for archival and forget the ones flagged for
252/// deletion, capped per run, the same way the maintenance sweep does.
253/// Memories flagged for archiving or consolidation are only counted.
254async fn maintenance_archive(db: Db) -> Result<Json<Value>, ApiError> {
255    let handle = db.handle.clone();
256    let cap = ops_config().max_forgets_per_archive;
257    let report = blocking("archive", move || {
258        let decisions = handle
259            .evaluate_archival_global()
260            .map_err(internal("archival"))?;
261        let evaluated = decisions.len();
262        let (mut keep, mut archive, mut consolidate, mut delete) = (0usize, 0usize, 0, 0);
263        let mut forgotten = 0usize;
264        for (id, decision) in decisions {
265            match decision {
266                ArchivalDecision::Keep => keep += 1,
267                ArchivalDecision::Archive => archive += 1,
268                ArchivalDecision::Consolidate(_) => consolidate += 1,
269                ArchivalDecision::Delete => {
270                    delete += 1;
271                    if forgotten < cap && handle.forget(id).is_ok() {
272                        forgotten += 1;
273                    }
274                }
275            }
276        }
277        handle.flush().map_err(internal("flush"))?;
278        Ok(json!({
279            "evaluated": evaluated,
280            "keep": keep,
281            "flagged_archive": archive,
282            "flagged_consolidate": consolidate,
283            "flagged_delete": delete,
284            "forgotten": forgotten,
285            "capped": delete > forgotten,
286        }))
287    })
288    .await?;
289    Ok(Json(with_database(&db.name, report)))
290}
291
292/// Merge duplicates: first the paraphrase dedup pass (the newer restatement
293/// is folded into the older memory), then the sweep's cluster consolidation
294/// (near duplicates merged into one memory, sources invalidated and linked).
295async fn maintenance_consolidate(db: Db) -> Result<Json<Value>, ApiError> {
296    let handle = db.handle.clone();
297    let cfg = ops_config();
298    let (batch, min_cluster, similarity) = (
299        cfg.dedup_batch_size.max(1),
300        cfg.consolidation_min_cluster_size,
301        cfg.consolidation_similarity,
302    );
303    let report = blocking("consolidate", move || {
304        let ids = handle.memory_ids();
305        let mut duplicates_merged = 0usize;
306        for chunk in ids.chunks(batch) {
307            duplicates_merged += handle.dedup_sweep_ids(chunk).map_err(internal("dedup"))?;
308        }
309        let candidates = handle
310            .find_consolidation_candidates(min_cluster, similarity)
311            .map_err(internal("consolidation"))?;
312        let clusters_found = candidates.len();
313        let clusters_consolidated = candidates
314            .iter()
315            .filter(|c| handle.consolidate_cluster(&c.memories).is_ok())
316            .count();
317        handle.flush().map_err(internal("flush"))?;
318        Ok(json!({
319            "duplicates_merged": duplicates_merged,
320            "clusters_found": clusters_found,
321            "clusters_consolidated": clusters_consolidated,
322        }))
323    })
324    .await?;
325    Ok(Json(with_database(&db.name, report)))
326}
327
328/// The full maintenance sweep the scheduler runs.
329async fn maintenance_sweep(db: Db) -> Result<Json<Value>, ApiError> {
330    let handle = db.handle.clone();
331    let report = blocking("sweep", move || Ok(crate::maintenance::run_sweep(&handle))).await?;
332    Ok(Json(json!({
333        "database": db.name,
334        "consolidated": report.consolidated,
335        "decayed": report.decayed,
336        "forgotten": report.forgotten,
337    })))
338}
339
340fn with_database(name: &str, mut report: Value) -> Value {
341    if let Value::Object(map) = &mut report {
342        map.insert("database".into(), Value::String(name.to_string()));
343    }
344    report
345}
346
347// ---------------------------------------------------------------------------
348// POST /enrich
349// ---------------------------------------------------------------------------
350
351/// Run the enrichment pipeline over this database's unenriched episodic
352/// memories, then entity linking, communities and the user model.
353async fn enrich(State(state): State<Arc<AppState>>, db: Db) -> Result<Json<Value>, ApiError> {
354    let config = state.extraction_config.clone().ok_or_else(|| {
355        ApiError::ServiceUnavailable(
356            "enrichment needs an LLM: start the server with --llm-provider or \
357             MENTEDB_LLM_PROVIDER"
358                .into(),
359        )
360    })?;
361    let embedder = state.databases.embedder().ok_or_else(|| {
362        ApiError::ServiceUnavailable(
363            "enrichment needs an embedder: start the server with an embedding provider".into(),
364        )
365    })?;
366    let provider = mentedb_extraction::HttpExtractionProvider::new(config.clone())
367        .map_err(|e| ApiError::ServiceUnavailable(format!("LLM provider unavailable: {e}")))?;
368    let llm = mentedb::cognitive::CognitiveLlmService::new(
369        mentedb_extraction::ExtractionLlmJudge::new(provider),
370    );
371    let handle = db.handle.clone();
372    let runtime = tokio::runtime::Handle::current();
373    let result = blocking("enrichment", move || {
374        // The engine compares this watermark with memory creation times, so
375        // the next run only considers memories stored after this one.
376        let watermark = now_micros();
377        let result = runtime.block_on(mentedb::enrichment::run_enrichment(
378            &handle,
379            config,
380            &*embedder,
381            Some(&llm),
382            watermark,
383            false,
384        ));
385        handle.flush().map_err(internal("flush"))?;
386        Ok(result)
387    })
388    .await?;
389    Ok(Json(json!({
390        "database": db.name,
391        "memories_stored": result.memories_stored,
392        "edges_created": result.edges_created,
393        "entities_extracted": result.entities_extracted,
394        "duplicates_skipped": result.duplicates_skipped,
395        "contradictions_found": result.contradictions_found,
396        "sync_linked": result.sync_linked,
397        "llm_linked": result.llm_linked,
398        "communities_created": result.communities_created,
399        "user_model_updated": result.user_model_updated,
400    })))
401}
402
403// ---------------------------------------------------------------------------
404// POST /documents, DELETE /documents/{document_id}
405// ---------------------------------------------------------------------------
406
407/// Tag carried by every chunk of a document.
408fn document_tag(document_id: &Uuid) -> String {
409    format!("doc:{document_id}")
410}
411
412fn optional_str<'a>(req: &'a Value, field: &str) -> Result<Option<&'a str>, ApiError> {
413    match req.get(field) {
414        None | Some(Value::Null) => Ok(None),
415        Some(Value::String(s)) => Ok(Some(s.as_str())),
416        Some(_) => Err(ApiError::BadRequest(format!("'{field}' must be a string"))),
417    }
418}
419
420fn optional_usize(req: &Value, field: &str) -> Result<Option<usize>, ApiError> {
421    match req.get(field) {
422        None | Some(Value::Null) => Ok(None),
423        Some(v) => v
424            .as_u64()
425            .map(|n| Some(n as usize))
426            .ok_or_else(|| ApiError::BadRequest(format!("'{field}' must be a whole number"))),
427    }
428}
429
430/// Load a document into memory as overlapping chunks.
431///
432/// Body: `content` (required), optional `title`, `source`, `agent_id`,
433/// `user_id`, `tags`, `memory_type` (default `semantic`), `chunk_chars`,
434/// `overlap_chars`.
435async fn store_document(
436    agent: Option<Extension<AuthenticatedAgent>>,
437    db: Db,
438    Json(req): Json<Value>,
439) -> Result<impl IntoResponse, ApiError> {
440    let cfg = &ops_config().documents;
441    let content = req
442        .get("content")
443        .and_then(|v| v.as_str())
444        .filter(|s| !s.trim().is_empty())
445        .ok_or_else(|| ApiError::BadRequest("missing or empty 'content'".into()))?;
446    let title = optional_str(&req, "title")?.map(String::from);
447    let source = optional_str(&req, "source")?.map(String::from);
448
449    let agent_id = match crate::handlers::parse_optional_uuid(&req, "agent_id")? {
450        Some(id) => {
451            let id = AgentId(id);
452            crate::handlers::check_body_agent(&agent, id)?;
453            id
454        }
455        None => match &agent {
456            Some(Extension(authed)) => authed.agent()?.unwrap_or_else(AgentId::nil),
457            None => AgentId::nil(),
458        },
459    };
460    let user_id = crate::handlers::parse_optional_uuid(&req, "user_id")?
461        .map(UserId)
462        .unwrap_or_else(UserId::nil);
463    let memory_type = if req.get("memory_type").is_some_and(|v| !v.is_null()) {
464        crate::handlers::parse_memory_type(&req)?
465    } else {
466        MemoryType::Semantic
467    };
468    let mut tags: Vec<String> = match req.get("tags") {
469        None | Some(Value::Null) => vec![],
470        Some(Value::Array(arr)) => arr
471            .iter()
472            .map(|v| {
473                v.as_str()
474                    .map(String::from)
475                    .ok_or_else(|| ApiError::BadRequest("'tags' must be strings".into()))
476            })
477            .collect::<Result<_, _>>()?,
478        Some(_) => return Err(ApiError::BadRequest("'tags' must be an array".into())),
479    };
480
481    let chunk_chars = optional_usize(&req, "chunk_chars")?.unwrap_or(cfg.chunk_chars);
482    let overlap_chars = optional_usize(&req, "overlap_chars")?.unwrap_or(cfg.overlap_chars);
483    if chunk_chars < cfg.min_chunk_chars || chunk_chars > cfg.max_chunk_chars {
484        return Err(ApiError::BadRequest(format!(
485            "'chunk_chars' must be between {} and {}",
486            cfg.min_chunk_chars, cfg.max_chunk_chars
487        )));
488    }
489    let max_overlap = (chunk_chars as f32 * cfg.max_overlap_ratio) as usize;
490    if overlap_chars > max_overlap {
491        return Err(ApiError::BadRequest(format!(
492            "'overlap_chars' must be at most {max_overlap} for chunks of {chunk_chars}"
493        )));
494    }
495
496    let chunks = chunk_text(
497        content,
498        chunk_chars,
499        overlap_chars,
500        &cfg.sentence_terminators,
501    );
502    if chunks.len() > cfg.max_chunks {
503        return Err(ApiError::BadRequest(format!(
504            "document splits into {} chunks, more than the limit of {}",
505            chunks.len(),
506            cfg.max_chunks
507        )));
508    }
509
510    let document_id = Uuid::new_v4();
511    let doc_tag = document_tag(&document_id);
512    if !tags.contains(&doc_tag) {
513        tags.push(doc_tag);
514    }
515    let handle = db.handle.clone();
516    let chunk_count = chunks.len();
517    let ids = blocking("store document", move || {
518        let texts: Vec<&str> = chunks.iter().map(String::as_str).collect();
519        let embeddings = handle
520            .embed_batch(&texts)
521            .map_err(|e| ApiError::ServiceUnavailable(format!("embedding failed: {e}")))?;
522        let now = now_micros();
523        let nodes: Vec<MemoryNode> = chunks
524            .iter()
525            .enumerate()
526            .map(|(i, text)| {
527                let mut attributes = HashMap::new();
528                attributes.insert(
529                    "document_id".to_string(),
530                    AttributeValue::String(document_id.to_string()),
531                );
532                attributes.insert("chunk_index".to_string(), AttributeValue::Integer(i as i64));
533                attributes.insert(
534                    "chunk_count".to_string(),
535                    AttributeValue::Integer(chunk_count as i64),
536                );
537                if let Some(t) = &title {
538                    attributes.insert("title".to_string(), AttributeValue::String(t.clone()));
539                }
540                if let Some(s) = &source {
541                    attributes.insert("source".to_string(), AttributeValue::String(s.clone()));
542                }
543                MemoryNode {
544                    id: MemoryId::new(),
545                    agent_id,
546                    user_id,
547                    memory_type,
548                    embedding: embeddings
549                        .as_ref()
550                        .and_then(|e| e.get(i).cloned())
551                        .unwrap_or_default(),
552                    content: text.clone(),
553                    created_at: now,
554                    accessed_at: now,
555                    access_count: 0,
556                    salience: 0.5,
557                    confidence: 1.0,
558                    space_id: SpaceId::nil(),
559                    attributes,
560                    tags: tags.clone(),
561                    valid_from: None,
562                    valid_until: None,
563                    context: None,
564                }
565            })
566            .collect();
567        let ids = handle.store_batch(nodes).map_err(internal("store"))?;
568        handle.flush().map_err(internal("flush"))?;
569        Ok(ids)
570    })
571    .await?;
572
573    Ok((
574        StatusCode::CREATED,
575        Json(json!({
576            "database": db.name,
577            "document_id": document_id.to_string(),
578            "chunk_count": ids.len(),
579            "chunk_ids": ids.iter().map(|id| id.to_string()).collect::<Vec<_>>(),
580        })),
581    ))
582}
583
584/// Forget every chunk of a document.
585async fn delete_document(
586    agent: Option<Extension<AuthenticatedAgent>>,
587    db: Db,
588    Path(params): Path<HashMap<String, String>>,
589) -> Result<Json<Value>, ApiError> {
590    let document_id: Uuid = params
591        .get("document_id")
592        .ok_or_else(|| ApiError::BadRequest("missing document id".into()))?
593        .parse()
594        .map_err(|_| ApiError::BadRequest("invalid document id".into()))?;
595    let handle = db.handle.clone();
596    let deleted = blocking("delete document", move || {
597        let ids = handle.memory_ids_with_tag(&document_tag(&document_id));
598        if ids.is_empty() {
599            return Err(ApiError::NotFound(format!(
600                "document {document_id} not found"
601            )));
602        }
603        // Check every chunk before forgetting any, so a refused request
604        // leaves the document whole.
605        if agent.is_some() {
606            for id in &ids {
607                if let Ok(node) = handle.get_memory(*id) {
608                    crate::handlers::check_owner(&agent, &node)?;
609                }
610            }
611        }
612        for id in &ids {
613            handle.forget(*id).map_err(internal("forget"))?;
614        }
615        handle.flush().map_err(internal("flush"))?;
616        Ok(ids.len())
617    })
618    .await?;
619    Ok(Json(json!({
620        "database": db.name,
621        "document_id": document_id.to_string(),
622        "deleted": deleted,
623    })))
624}
625
626/// Split text into chunks of at most `chunk_chars` characters. Paragraphs
627/// (separated by blank lines) are kept whole when they fit; longer ones are
628/// split on sentence boundaries, and sentences longer than a chunk on word
629/// boundaries. Each chunk after the first starts with up to `overlap_chars`
630/// characters from the end of the previous one, cut at a word boundary.
631pub fn chunk_text(
632    text: &str,
633    chunk_chars: usize,
634    overlap_chars: usize,
635    terminators: &[char],
636) -> Vec<String> {
637    let chunk_chars = chunk_chars.max(1);
638    let overlap_chars = overlap_chars.min(chunk_chars / 2);
639    // Room left for new text once the overlap and a separator are in place.
640    let max_unit = chunk_chars.saturating_sub(overlap_chars + 2).max(1);
641
642    // (unit, starts a new paragraph)
643    let mut units: Vec<(String, bool)> = Vec::new();
644    for paragraph in paragraphs(text) {
645        if char_len(&paragraph) <= max_unit {
646            units.push((paragraph, true));
647            continue;
648        }
649        let mut first = true;
650        for sentence in sentences(&paragraph, terminators) {
651            for piece in split_long(&sentence, max_unit) {
652                units.push((piece, first));
653                first = false;
654            }
655        }
656    }
657
658    let mut chunks = Vec::new();
659    let mut current = String::new();
660    for (unit, new_paragraph) in units {
661        let sep = if new_paragraph { "\n\n" } else { " " };
662        if current.is_empty() {
663            current = unit;
664        } else if char_len(&current) + sep.len() + char_len(&unit) <= chunk_chars {
665            current.push_str(sep);
666            current.push_str(&unit);
667        } else {
668            let tail = overlap_tail(&current, overlap_chars);
669            chunks.push(std::mem::take(&mut current));
670            if !tail.is_empty() {
671                current.push_str(&tail);
672                current.push_str(sep);
673            }
674            current.push_str(&unit);
675        }
676    }
677    if !current.is_empty() {
678        chunks.push(current);
679    }
680    chunks
681}
682
683fn char_len(s: &str) -> usize {
684    s.chars().count()
685}
686
687/// Paragraphs separated by blank lines, trimmed, empty ones dropped.
688fn paragraphs(text: &str) -> Vec<String> {
689    let mut out = Vec::new();
690    let mut lines: Vec<&str> = Vec::new();
691    for line in text.lines() {
692        if line.trim().is_empty() {
693            if !lines.is_empty() {
694                out.push(lines.join("\n").trim().to_string());
695                lines.clear();
696            }
697        } else {
698            lines.push(line);
699        }
700    }
701    if !lines.is_empty() {
702        out.push(lines.join("\n").trim().to_string());
703    }
704    out.retain(|p| !p.is_empty());
705    out
706}
707
708/// Sentences: a terminator followed by whitespace ends one.
709fn sentences(paragraph: &str, terminators: &[char]) -> Vec<String> {
710    let mut out = Vec::new();
711    let mut current = String::new();
712    let mut chars = paragraph.chars().peekable();
713    while let Some(c) = chars.next() {
714        current.push(c);
715        if terminators.contains(&c) && chars.peek().is_none_or(|n| n.is_whitespace()) {
716            let s = current.trim();
717            if !s.is_empty() {
718                out.push(s.to_string());
719            }
720            current.clear();
721        }
722    }
723    let s = current.trim();
724    if !s.is_empty() {
725        out.push(s.to_string());
726    }
727    out
728}
729
730/// Pieces of at most `max` characters, broken at whitespace when possible.
731fn split_long(text: &str, max: usize) -> Vec<String> {
732    let mut out = Vec::new();
733    let mut current = String::new();
734    for word in text.split_whitespace() {
735        let mut word = word.to_string();
736        // A single word longer than a piece is cut by characters.
737        while char_len(&word) > max {
738            if !current.is_empty() {
739                out.push(std::mem::take(&mut current));
740            }
741            let head: String = word.chars().take(max).collect();
742            word = word.chars().skip(max).collect();
743            out.push(head);
744        }
745        if word.is_empty() {
746            continue;
747        }
748        if current.is_empty() {
749            current = word;
750        } else if char_len(&current) + 1 + char_len(&word) <= max {
751            current.push(' ');
752            current.push_str(&word);
753        } else {
754            out.push(std::mem::replace(&mut current, word));
755        }
756    }
757    if !current.is_empty() {
758        out.push(current);
759    }
760    out
761}
762
763/// The last `overlap` characters of `chunk`, starting at a word boundary.
764fn overlap_tail(chunk: &str, overlap: usize) -> String {
765    if overlap == 0 {
766        return String::new();
767    }
768    let total = char_len(chunk);
769    if total <= overlap {
770        return chunk.trim().to_string();
771    }
772    let start = chunk
773        .char_indices()
774        .nth(total - overlap)
775        .map(|(i, _)| i)
776        .unwrap_or(0);
777    let tail = &chunk[start..];
778    // Drop a partial leading word unless the tail is one single word.
779    let preceded_by_space = chunk[..start].ends_with(char::is_whitespace);
780    let tail = if preceded_by_space {
781        tail
782    } else {
783        match tail.find(char::is_whitespace) {
784            Some(i) => &tail[i..],
785            None => tail,
786        }
787    };
788    tail.trim().to_string()
789}
790
791// ---------------------------------------------------------------------------
792// POST /backup, GET /backups
793// ---------------------------------------------------------------------------
794
795/// File extension the storage layer lists as a backup.
796const BACKUP_EXTENSION: &str = "mentebackup";
797
798fn backup_root(state: &AppState) -> PathBuf {
799    ops_config()
800        .backup_dir
801        .clone()
802        .unwrap_or_else(|| state.databases.root().join("backups"))
803}
804
805/// `YYYYMMDDTHHMMSS.mmmZ` in UTC.
806fn utc_stamp(at: SystemTime) -> String {
807    let since = at.duration_since(UNIX_EPOCH).unwrap_or_default();
808    let secs = since.as_secs();
809    let (days, rem) = ((secs / 86_400) as i64, secs % 86_400);
810    let (year, month, day) = civil_from_days(days);
811    format!(
812        "{year:04}{month:02}{day:02}T{:02}{:02}{:02}.{:03}Z",
813        rem / 3600,
814        (rem % 3600) / 60,
815        rem % 60,
816        since.subsec_millis()
817    )
818}
819
820/// Proleptic Gregorian date for days since 1970-01-01 (Howard Hinnant's
821/// civil_from_days).
822fn civil_from_days(days: i64) -> (i64, u32, u32) {
823    let z = days + 719_468;
824    let era = z.div_euclid(146_097);
825    let doe = z.rem_euclid(146_097);
826    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
827    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
828    let mp = (5 * doy + 2) / 153;
829    let day = (doy - (153 * mp + 2) / 5 + 1) as u32;
830    let month = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
831    let year = yoe + era * 400 + i64::from(month <= 2);
832    (year, month, day)
833}
834
835/// Copy `src` into `dst`, skipping any path under one of `exclude`.
836fn copy_tree(src: &FsPath, dst: &FsPath, exclude: &[PathBuf]) -> std::io::Result<()> {
837    std::fs::create_dir_all(dst)?;
838    for entry in std::fs::read_dir(src)? {
839        let entry = entry?;
840        let path = entry.path();
841        let canonical = path.canonicalize().unwrap_or_else(|_| path.clone());
842        if exclude.iter().any(|x| canonical.starts_with(x)) {
843            continue;
844        }
845        let target = dst.join(entry.file_name());
846        if entry.file_type()?.is_dir() {
847            copy_tree(&path, &target, exclude)?;
848        } else {
849            std::fs::copy(&path, &target)?;
850        }
851    }
852    Ok(())
853}
854
855/// Flush the database and write a full backup file to
856/// `<backup_dir>/<database>/<database>-<UTC time>.mentebackup`.
857///
858/// Writes that land while the files are read may or may not be included;
859/// for a snapshot at an exact point, pause writers first.
860async fn create_backup(
861    _admin: AdminGuard,
862    State(state): State<Arc<AppState>>,
863    db: Db,
864) -> Result<impl IntoResponse, ApiError> {
865    let db_dir = state.databases.path_of(&db.name)?;
866    let root = backup_root(&state);
867    let data_root = state.databases.root().to_path_buf();
868    let handle = db.handle.clone();
869    let name = db.name.clone();
870    let (file, manifest) = blocking("backup", move || {
871        handle.flush_full().map_err(internal("flush"))?;
872        let io = |what: &'static str| {
873            move |e: std::io::Error| ApiError::Internal(format!("{what}: {e}"))
874        };
875        let dir = root.join(&name);
876        std::fs::create_dir_all(&dir).map_err(io("create backup directory"))?;
877        let root_c = root
878            .canonicalize()
879            .map_err(io("resolve backup directory"))?;
880        let db_c = db_dir
881            .canonicalize()
882            .map_err(io("resolve database directory"))?;
883        if db_c.starts_with(&root_c) {
884            return Err(ApiError::Internal(
885                "the backup directory must not contain the database directory".into(),
886            ));
887        }
888        let file = format!("{name}-{}.{BACKUP_EXTENSION}", utc_stamp(SystemTime::now()));
889        let path = dir.join(&file);
890        if path.exists() {
891            return Err(ApiError::Conflict(format!(
892                "backup {file} already exists; retry"
893            )));
894        }
895
896        // The default database's directory also holds the named databases and,
897        // by default, the backups; those are left out of its backup.
898        let mut exclude = vec![root_c];
899        if name == DEFAULT_DB {
900            let dbs = data_root.join("dbs");
901            exclude.push(dbs.canonicalize().unwrap_or(dbs));
902        }
903        exclude.retain(|x| x.starts_with(&db_c));
904
905        let manifest = if exclude.is_empty() {
906            BackupManager::create_backup(&db_dir, &path).map_err(internal("backup"))?
907        } else {
908            let staging = dir.join(format!(".staging-{}", Uuid::new_v4()));
909            let staged = copy_tree(&db_dir, &staging, &exclude)
910                .map_err(io("stage backup"))
911                .and_then(|_| {
912                    BackupManager::create_backup(&staging, &path).map_err(internal("backup"))
913                });
914            let _ = std::fs::remove_dir_all(&staging);
915            staged?
916        };
917        Ok((file, manifest))
918    })
919    .await?;
920    Ok((
921        StatusCode::CREATED,
922        Json(json!({
923            "database": db.name,
924            "file": file,
925            "manifest": manifest,
926        })),
927    ))
928}
929
930/// Backups of this database, newest first.
931async fn list_backups(
932    _admin: AdminGuard,
933    State(state): State<Arc<AppState>>,
934    db: Db,
935) -> Result<Json<Value>, ApiError> {
936    let dir = backup_root(&state).join(&db.name);
937    let backups = blocking("list backups", move || {
938        let mut found = Vec::new();
939        let entries = match std::fs::read_dir(&dir) {
940            Ok(e) => e,
941            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(found),
942            Err(e) => return Err(ApiError::Internal(format!("list backups: {e}"))),
943        };
944        for entry in entries.flatten() {
945            let path = entry.path();
946            if path.extension().and_then(|e| e.to_str()) != Some(BACKUP_EXTENSION) {
947                continue;
948            }
949            if let Ok(manifest) = BackupManager::read_manifest(&path) {
950                found.push((entry.file_name().to_string_lossy().into_owned(), manifest));
951            }
952        }
953        found.sort_by_key(|(_, m)| std::cmp::Reverse(m.created_at));
954        Ok(found)
955    })
956    .await?;
957    let backups: Vec<Value> = backups
958        .into_iter()
959        .map(|(file, manifest)| json!({ "file": file, "manifest": manifest }))
960        .collect();
961    Ok(Json(json!({ "database": db.name, "backups": backups })))
962}
963
964// ---------------------------------------------------------------------------
965// /v1/databases
966// ---------------------------------------------------------------------------
967
968/// The default database and every named one on disk.
969async fn list_databases(
970    _admin: AdminGuard,
971    State(state): State<Arc<AppState>>,
972) -> Result<Json<Value>, ApiError> {
973    let dbs = state.databases.clone();
974    let (names, open) = blocking("list databases", move || {
975        let open: Vec<String> = dbs.open_handles().into_iter().map(|(n, _)| n).collect();
976        Ok((dbs.list(), open))
977    })
978    .await?;
979    Ok(Json(json!({
980        "default": DEFAULT_DB,
981        "databases": names,
982        "open": open,
983    })))
984}
985
986/// Create a named database. Body: `{"name": "..."}`.
987async fn create_database(
988    _admin: AdminGuard,
989    State(state): State<Arc<AppState>>,
990    Json(req): Json<Value>,
991) -> Result<impl IntoResponse, ApiError> {
992    let name = req
993        .get("name")
994        .and_then(|v| v.as_str())
995        .ok_or_else(|| ApiError::BadRequest("missing 'name'".into()))?
996        .to_string();
997    validate_name(&name)?;
998    let dbs = state.databases.clone();
999    let created = name.clone();
1000    blocking("create database", move || {
1001        if dbs.exists(&created) {
1002            return Err(ApiError::Conflict(format!(
1003                "database '{created}' already exists"
1004            )));
1005        }
1006        dbs.get(&created, true).map(|_| ())
1007    })
1008    .await?;
1009    Ok((
1010        StatusCode::CREATED,
1011        Json(json!({ "name": name, "status": "created" })),
1012    ))
1013}
1014
1015/// Close and permanently delete a named database.
1016async fn delete_database(
1017    _admin: AdminGuard,
1018    State(state): State<Arc<AppState>>,
1019    Path(name): Path<String>,
1020) -> Result<Json<Value>, ApiError> {
1021    validate_name(&name)?;
1022    let dbs = state.databases.clone();
1023    let target = name.clone();
1024    blocking("delete database", move || dbs.delete(&target)).await?;
1025    Ok(Json(json!({ "name": name, "status": "deleted" })))
1026}
1027
1028// ---------------------------------------------------------------------------
1029// GET /v1/ready
1030// ---------------------------------------------------------------------------
1031
1032/// The last embed check for one embedder. The `Weak` keeps the embedder's
1033/// allocation from being reused, so a pointer match is always the same one.
1034struct EmbedProbe {
1035    embedder: Weak<dyn EmbeddingProvider>,
1036    checked_at: Instant,
1037    result: Result<(), String>,
1038}
1039
1040static EMBED_PROBES: Mutex<Vec<EmbedProbe>> = Mutex::new(Vec::new());
1041
1042async fn check_embedder(
1043    embedder: Arc<dyn EmbeddingProvider>,
1044    cfg: &ReadinessConfig,
1045) -> Result<(), String> {
1046    let weak = Arc::downgrade(&embedder);
1047    {
1048        let mut probes = EMBED_PROBES.lock().unwrap_or_else(|e| e.into_inner());
1049        probes.retain(|p| p.embedder.strong_count() > 0);
1050        if let Some(p) = probes.iter().find(|p| Weak::ptr_eq(&p.embedder, &weak))
1051            && p.checked_at.elapsed() < cfg.embed_check_interval
1052        {
1053            return p.result.clone();
1054        }
1055    }
1056    let text = cfg.probe_text.clone();
1057    let probe = tokio::task::spawn_blocking(move || embedder.embed(&text));
1058    let result = match tokio::time::timeout(cfg.check_timeout, probe).await {
1059        Ok(Ok(Ok(v))) if !v.is_empty() => Ok(()),
1060        Ok(Ok(Ok(_))) => Err("embedder returned an empty vector".to_string()),
1061        Ok(Ok(Err(e))) => Err(format!("embed failed: {e}")),
1062        Ok(Err(e)) => Err(format!("embed task failed: {e}")),
1063        Err(_) => Err(format!("embed timed out after {:?}", cfg.check_timeout)),
1064    };
1065    let mut probes = EMBED_PROBES.lock().unwrap_or_else(|e| e.into_inner());
1066    probes.retain(|p| !Weak::ptr_eq(&p.embedder, &weak));
1067    probes.push(EmbedProbe {
1068        embedder: weak,
1069        checked_at: Instant::now(),
1070        result: result.clone(),
1071    });
1072    result
1073}
1074
1075/// Readiness: 200 when the default database answers and, if an embedder is
1076/// configured, it embeds (checked at most once per configured interval).
1077/// 503 naming the failing check otherwise.
1078async fn ready(State(state): State<Arc<AppState>>) -> Response {
1079    let cfg = &ops_config().readiness;
1080    let mut checks = serde_json::Map::new();
1081    let mut failed: Vec<&str> = Vec::new();
1082
1083    let db = state.databases.default_db();
1084    let probe = tokio::task::spawn_blocking(move || db.memory_count());
1085    match tokio::time::timeout(cfg.check_timeout, probe).await {
1086        Ok(Ok(_)) => {
1087            checks.insert("database".into(), json!("ok"));
1088        }
1089        Ok(Err(e)) => {
1090            checks.insert("database".into(), json!(format!("error: {e}")));
1091            failed.push("database");
1092        }
1093        Err(_) => {
1094            checks.insert(
1095                "database".into(),
1096                json!(format!("error: timed out after {:?}", cfg.check_timeout)),
1097            );
1098            failed.push("database");
1099        }
1100    }
1101
1102    match state.databases.embedder() {
1103        None => {
1104            checks.insert("embedder".into(), json!("not_configured"));
1105        }
1106        Some(embedder) => match check_embedder(embedder, cfg).await {
1107            Ok(()) => {
1108                checks.insert("embedder".into(), json!("ok"));
1109            }
1110            Err(e) => {
1111                checks.insert("embedder".into(), json!(format!("error: {e}")));
1112                failed.push("embedder");
1113            }
1114        },
1115    }
1116
1117    if failed.is_empty() {
1118        (
1119            StatusCode::OK,
1120            Json(json!({ "status": "ready", "checks": checks })),
1121        )
1122            .into_response()
1123    } else {
1124        (
1125            StatusCode::SERVICE_UNAVAILABLE,
1126            Json(json!({ "status": "not_ready", "failed": failed, "checks": checks })),
1127        )
1128            .into_response()
1129    }
1130}
1131
1132#[cfg(test)]
1133mod tests {
1134    use super::*;
1135
1136    const TERMS: &[char] = &['.', '!', '?'];
1137
1138    #[test]
1139    fn short_text_is_one_chunk() {
1140        let chunks = chunk_text("One paragraph. Two sentences.", 100, 20, TERMS);
1141        assert_eq!(chunks, vec!["One paragraph. Two sentences.".to_string()]);
1142    }
1143
1144    #[test]
1145    fn chunks_respect_the_limit_and_overlap() {
1146        let sentence = "The quick brown fox jumps over the lazy dog again.";
1147        let paragraph = [sentence; 6].join(" ");
1148        let text = [paragraph.as_str(); 4].join("\n\n");
1149        let chunks = chunk_text(&text, 200, 40, TERMS);
1150        assert!(chunks.len() > 4);
1151        for c in &chunks {
1152            assert!(char_len(c) <= 200, "{} chars: {c}", char_len(c));
1153        }
1154        for pair in chunks.windows(2) {
1155            let tail_word = pair[0].split_whitespace().last().unwrap();
1156            assert!(pair[1].contains(tail_word));
1157        }
1158    }
1159
1160    #[test]
1161    fn a_word_longer_than_a_chunk_is_cut() {
1162        let text = "x".repeat(450);
1163        let chunks = chunk_text(&text, 200, 0, TERMS);
1164        assert!(chunks.iter().all(|c| char_len(c) <= 200));
1165        assert_eq!(chunks.concat(), text);
1166    }
1167
1168    #[test]
1169    fn utc_stamp_formats_known_instants() {
1170        assert_eq!(utc_stamp(UNIX_EPOCH), "19700101T000000.000Z");
1171        let t = UNIX_EPOCH + Duration::from_millis(1_709_210_096_789);
1172        assert_eq!(utc_stamp(t), "20240229T123456.789Z");
1173    }
1174}