miku/
db.rs

1use anyhow::Result;
2use libsql::{Builder, Connection, Database};
3use std::collections::HashMap;
4use std::env;
5
6pub const DEFAULT_SEARCH_LIMIT: usize = 100;
7pub use crate::constant::ENV_DATABASE_URL;
8
9pub async fn init_db() -> Result<(Database, Connection)> {
10    let paths = crate::paths::MikuPaths::resolve();
11    if !paths.data_dir.exists() {
12        std::fs::create_dir_all(&paths.data_dir)?;
13    }
14
15    let db_path = env::var(ENV_DATABASE_URL)
16        .unwrap_or_else(|_| paths.db_path().to_str().unwrap().to_string());
17    let db = Builder::new_local(&db_path).build().await?;
18    let conn = db.connect()?;
19
20    conn.execute(crate::constant::PRAGMA_FOREIGN_KEYS_ON, ())
21        .await?;
22
23    conn.execute(crate::constant::SCHEMA_CREATE_TOPICS, ())
24        .await?;
25
26    // FTS5 for full-text keyword search
27    conn.execute(crate::constant::SCHEMA_CREATE_TOPICS_FTS, ())
28        .await?;
29
30    conn.execute(crate::constant::SCHEMA_CREATE_SESSIONS, ())
31        .await?;
32
33    // Graph Tier (Hot)
34    conn.execute(crate::constant::SCHEMA_CREATE_MCP_ENTITIES, ())
35        .await?;
36
37    conn.execute(crate::constant::SCHEMA_CREATE_MCP_OBSERVATIONS, ())
38        .await?;
39
40    conn.execute(crate::constant::SCHEMA_CREATE_IDX_MCP_OBSERVATIONS, ())
41        .await?;
42
43    conn.execute(crate::constant::SCHEMA_CREATE_MCP_RELATIONS, ())
44        .await?;
45
46    conn.execute(crate::constant::SCHEMA_CREATE_MCP_TRUTHS, ())
47        .await?;
48
49    conn.execute(crate::constant::SCHEMA_CREATE_MCP_SKILLS, ())
50        .await?;
51
52    // Document Tier (Vectors)
53    conn.execute(crate::constant::SCHEMA_CREATE_CHUNKS, ())
54        .await?;
55
56    conn.execute(crate::constant::SCHEMA_CREATE_IDX_CHUNKS_TOPIC_ID, ())
57        .await?;
58
59    // Vector index - metric=cosine is default
60    conn.execute(crate::constant::SCHEMA_CREATE_IDX_CHUNKS_VECTOR, ())
61        .await?;
62
63    // Triggers to keep FTS5 in sync with topics
64    conn.execute(crate::constant::SCHEMA_CREATE_TRIGGER_TOPICS_AI, ())
65        .await?;
66    conn.execute(crate::constant::SCHEMA_CREATE_TRIGGER_TOPICS_AD, ())
67        .await?;
68    conn.execute(crate::constant::SCHEMA_CREATE_TRIGGER_TOPICS_AU, ())
69        .await?;
70
71    // FTS5 for graph observation search (porter stemming, BM25 ranking)
72    conn.execute(crate::constant::SCHEMA_CREATE_MCP_OBS_FTS, ())
73        .await?;
74
75    // Triggers to keep mcp_obs_fts in sync with mcp_observations
76    conn.execute(crate::constant::SCHEMA_CREATE_TRIGGER_MCP_OBS_AI, ())
77        .await?;
78    conn.execute(crate::constant::SCHEMA_CREATE_TRIGGER_MCP_OBS_AD, ())
79        .await?;
80    conn.execute(crate::constant::SCHEMA_CREATE_TRIGGER_MCP_OBS_AU, ())
81        .await?;
82
83    Ok((db, conn))
84}
85
86/// FTS5 keyword search — returns (id, title, file_path, bm25_score)
87pub async fn search_fts(
88    conn: &Connection,
89    query: &str,
90    limit: usize,
91) -> Result<Vec<(String, String, String, f64)>> {
92    let mut rows = conn
93        .query(
94            crate::constant::SQL_SEARCH_FTS,
95            libsql::params![query, limit as i64],
96        )
97        .await?;
98    let mut results = Vec::new();
99    while let Some(row) = rows.next().await? {
100        results.push((
101            row.get::<String>(0)?,
102            row.get::<String>(1)?,
103            row.get::<String>(2)?,
104            row.get::<f64>(3)?,
105        ));
106    }
107    Ok(results)
108}
109
110pub async fn upsert_topic(
111    conn: &Connection,
112    id: &str,
113    title: &str,
114    file_path: &str,
115    body: &str,
116) -> Result<()> {
117    conn.execute(
118        crate::constant::SQL_UPSERT_TOPIC,
119        libsql::params![id, title, file_path, body],
120    )
121    .await?;
122    Ok(())
123}
124
125pub async fn delete_topic(conn: &Connection, id: &str) -> Result<()> {
126    conn.execute("DELETE FROM topics WHERE id = ?1", libsql::params![id])
127        .await?;
128    Ok(())
129}
130
131pub async fn mcp_create_entities(
132    conn: &Connection,
133    entities: Vec<crate::mcp::EntityInput>,
134) -> Result<()> {
135    let tx = conn.transaction().await?;
136    for mut ent in entities {
137        ent.name = crate::normalize::normalize_key(&ent.name);
138        tx.execute(
139            crate::constant::SQL_INSERT_ENTITY,
140            libsql::params![ent.name.clone(), ent.entity_type],
141        )
142        .await?;
143        for obs in ent.observations {
144            tx.execute(
145                crate::constant::SQL_INSERT_OBSERVATION,
146                libsql::params![uuid::Uuid::new_v4().to_string(), ent.name.clone(), obs],
147            )
148            .await?;
149        }
150    }
151    tx.commit().await?;
152    Ok(())
153}
154
155pub async fn mcp_add_observations(
156    conn: &Connection,
157    observations: Vec<crate::mcp::ObservationInput>,
158    limit: usize,
159) -> Result<()> {
160    let tx = conn.transaction().await?;
161    for mut obs_batch in observations {
162        obs_batch.entity_name = crate::normalize::normalize_key(&obs_batch.entity_name);
163        for content in obs_batch.contents {
164            tx.execute(
165                crate::constant::SQL_INSERT_OBSERVATION,
166                libsql::params![
167                    uuid::Uuid::new_v4().to_string(),
168                    obs_batch.entity_name.clone(),
169                    content
170                ],
171            )
172            .await?;
173        }
174        if limit > 0 {
175            tx.execute(
176                crate::constant::SQL_EVICT_OBSERVATIONS,
177                libsql::params![obs_batch.entity_name.clone(), limit as i64],
178            )
179            .await?;
180        }
181    }
182    tx.commit().await?;
183    Ok(())
184}
185
186pub async fn mcp_create_relations(
187    conn: &Connection,
188    relations: Vec<crate::mcp::RelationInput>,
189) -> Result<()> {
190    let tx = conn.transaction().await?;
191    for mut rel in relations {
192        rel.from = crate::normalize::normalize_key(&rel.from);
193        rel.to = crate::normalize::normalize_key(&rel.to);
194        tx.execute(
195            crate::constant::SQL_INSERT_RELATION,
196            libsql::params![rel.from, rel.to, rel.relation_type],
197        )
198        .await?;
199    }
200    tx.commit().await?;
201    Ok(())
202}
203
204pub async fn mcp_delete_entities(conn: &Connection, names: Vec<String>) -> Result<()> {
205    let tx = conn.transaction().await?;
206    for name in names {
207        let norm_name = crate::normalize::normalize_key(&name);
208        tx.execute(
209            crate::constant::SQL_DELETE_ENTITY,
210            libsql::params![norm_name.clone()],
211        )
212        .await?;
213        tx.execute(
214            "DELETE FROM topics WHERE id = ?1",
215            libsql::params![norm_name.clone()],
216        )
217        .await?;
218        tx.execute(
219            "DELETE FROM chunks WHERE topic_id = ?1",
220            libsql::params![norm_name],
221        )
222        .await?;
223    }
224    tx.commit().await?;
225    Ok(())
226}
227
228pub async fn mcp_delete_observations(
229    conn: &Connection,
230    deletions: Vec<crate::mcp::ObservationDeletion>,
231) -> Result<()> {
232    let tx = conn.transaction().await?;
233    for mut del in deletions {
234        del.entity_name = crate::normalize::normalize_key(&del.entity_name);
235        for obs in del.observations {
236            tx.execute(
237                crate::constant::SQL_DELETE_OBSERVATION,
238                libsql::params![del.entity_name.clone(), obs],
239            )
240            .await?;
241        }
242    }
243    tx.commit().await?;
244    Ok(())
245}
246
247pub async fn mcp_delete_relations(
248    conn: &Connection,
249    relations: Vec<crate::mcp::RelationInput>,
250) -> Result<()> {
251    let tx = conn.transaction().await?;
252    for mut rel in relations {
253        rel.from = crate::normalize::normalize_key(&rel.from);
254        rel.to = crate::normalize::normalize_key(&rel.to);
255        tx.execute(
256            crate::constant::SQL_DELETE_RELATION,
257            libsql::params![rel.from, rel.to, rel.relation_type],
258        )
259        .await?;
260    }
261    tx.commit().await?;
262    Ok(())
263}
264
265pub async fn mcp_read_graph(conn: &Connection) -> Result<crate::mcp::Graph> {
266    let mut entity_names = Vec::new();
267    let mut rows = conn
268        .query(crate::constant::SQL_SELECT_ALL_ENTITIES, ())
269        .await?;
270    while let Some(row) = rows.next().await? {
271        entity_names.push(row.get::<String>(0)?);
272    }
273    let entities = load_entities_lazy(conn, &entity_names).await?;
274
275    let mut relations = Vec::new();
276    let mut rel_rows = conn
277        .query(crate::constant::SQL_SELECT_ALL_RELATIONS, ())
278        .await?;
279    while let Some(row) = rel_rows.next().await? {
280        relations.push(crate::mcp::RelationInput {
281            from: row.get(0)?,
282            to: row.get(1)?,
283            relation_type: row.get(2)?,
284        });
285    }
286
287    Ok(crate::mcp::Graph {
288        entities,
289        relations,
290    })
291}
292
293pub async fn mcp_read_graph_eager(conn: &Connection) -> Result<crate::mcp::Graph> {
294    let mut entity_names = Vec::new();
295    let mut rows = conn
296        .query(crate::constant::SQL_SELECT_ALL_ENTITIES, ())
297        .await?;
298    while let Some(row) = rows.next().await? {
299        entity_names.push(row.get::<String>(0)?);
300    }
301    let entities = load_entities_eager(conn, &entity_names).await?;
302
303    let mut relations = Vec::new();
304    let mut rel_rows = conn
305        .query(crate::constant::SQL_SELECT_ALL_RELATIONS, ())
306        .await?;
307    while let Some(row) = rel_rows.next().await? {
308        relations.push(crate::mcp::RelationInput {
309            from: row.get(0)?,
310            to: row.get(1)?,
311            relation_type: row.get(2)?,
312        });
313    }
314
315    Ok(crate::mcp::Graph {
316        entities,
317        relations,
318    })
319}
320
321pub async fn mcp_search_nodes(conn: &Connection, query: &str) -> Result<crate::mcp::Graph> {
322    mcp_search_nodes_with_limit(conn, query, DEFAULT_SEARCH_LIMIT).await
323}
324
325pub async fn mcp_search_nodes_with_limit(
326    conn: &Connection,
327    query: &str,
328    limit: usize,
329) -> Result<crate::mcp::Graph> {
330    let limit = limit.max(1);
331    let mut entity_names: Vec<String> = Vec::new();
332
333    // Primary: FTS5 on observation content — porter stemming + BM25 ranking.
334    // Wrapped in an async block so any error (invalid syntax, bad token) is
335    // caught at the boundary and we fall through to the LIKE path.
336    let fts_hits: Vec<String> = async {
337        let fts_fetch_limit = limit.saturating_mul(8).max(limit) as i64;
338        let mut rows = conn
339            .query(
340                crate::constant::SQL_SEARCH_OBSERVATIONS_FTS,
341                libsql::params![query, fts_fetch_limit],
342            )
343            .await?;
344        let mut names = Vec::new();
345        while let Some(row) = rows.next().await? {
346            names.push(row.get::<String>(0)?);
347        }
348        Ok::<Vec<String>, anyhow::Error>(names)
349    }
350    .await
351    .unwrap_or_default();
352    for name in fts_hits {
353        if !entity_names.contains(&name) {
354            entity_names.push(name);
355            if entity_names.len() >= limit {
356                break;
357            }
358        }
359    }
360
361    // Secondary: LIKE on entity name / type — always runs, catches exact-name
362    // lookups and entity types that aren't in observations.
363    let pattern = format!("%{}%", query);
364    let mut rows = conn
365        .query(
366            crate::constant::SQL_SEARCH_ENTITIES_LIKE,
367            libsql::params![pattern, limit as i64],
368        )
369        .await?;
370    while let Some(row) = rows.next().await? {
371        let name: String = row.get(0)?;
372        if !entity_names.contains(&name) {
373            entity_names.push(name);
374            if entity_names.len() >= limit {
375                break;
376            }
377        }
378    }
379
380    // Expand neighbors (1-hop)
381    let relations = load_relations(conn, &entity_names).await?;
382    let mut all_entity_names = entity_names.clone();
383    for rel in &relations {
384        if !all_entity_names.contains(&rel.from) {
385            all_entity_names.push(rel.from.clone());
386        }
387        if !all_entity_names.contains(&rel.to) {
388            all_entity_names.push(rel.to.clone());
389        }
390    }
391
392    let entities = load_entities_lazy(conn, &all_entity_names).await?;
393
394    Ok(crate::mcp::Graph {
395        entities,
396        relations,
397    })
398}
399
400pub async fn mcp_open_nodes(conn: &Connection, names: Vec<String>) -> Result<crate::mcp::Graph> {
401    let normalized_names: Vec<String> = names
402        .into_iter()
403        .map(|n| crate::normalize::normalize_key(&n))
404        .collect();
405    let entities = load_entities_eager(conn, &normalized_names).await?;
406    let relations = load_relations(conn, &normalized_names).await?;
407
408    Ok(crate::mcp::Graph {
409        entities,
410        relations,
411    })
412}
413
414async fn load_relations(
415    conn: &Connection,
416    names: &[String],
417) -> Result<Vec<crate::mcp::RelationInput>> {
418    let mut relations = Vec::new();
419    if names.is_empty() {
420        return Ok(relations);
421    }
422
423    for chunk in names.chunks(400) {
424        let placeholders = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(",");
425        let sql = crate::constant::SQL_SELECT_RELATIONS_IN_TEMPLATE.replace("{0}", &placeholders);
426
427        let mut params = Vec::new();
428        // Since we use the same array twice in the query logic, we only pass params once, wait!
429        // The query "from_entity IN ({0}) OR to_entity IN ({0})" uses the placeholders twice.
430        // It's technically better to use different placeholders, but libsql bindings map "?" sequentially.
431        // So we need to push the parameters twice!
432        for name in chunk {
433            params.push(libsql::Value::from(name.clone()));
434        }
435        for name in chunk {
436            params.push(libsql::Value::from(name.clone()));
437        }
438
439        let mut rel_rows = conn.query(&sql, params).await?;
440        while let Some(row) = rel_rows.next().await? {
441            relations.push(crate::mcp::RelationInput {
442                from: row.get(0)?,
443                to: row.get(1)?,
444                relation_type: row.get(2)?,
445            });
446        }
447    }
448
449    Ok(relations)
450}
451
452async fn load_entities_lazy(
453    conn: &Connection,
454    names: &[String],
455) -> Result<Vec<crate::mcp::EntityOutput>> {
456    if names.is_empty() {
457        return Ok(Vec::new());
458    }
459
460    let mut entity_types = HashMap::new();
461    let mut obs_counts: HashMap<String, usize> = HashMap::new();
462    let truths = select_truths(conn, names).await?;
463
464    for chunk in names.chunks(500) {
465        let placeholders = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(",");
466
467        // Load entities
468        let entity_sql =
469            crate::constant::SQL_SELECT_ENTITIES_IN_TEMPLATE.replace("{}", &placeholders);
470        let params = chunk
471            .iter()
472            .cloned()
473            .map(libsql::Value::from)
474            .collect::<Vec<_>>();
475
476        let mut rows = conn.query(&entity_sql, params.clone()).await?;
477        while let Some(row) = rows.next().await? {
478            entity_types.insert(row.get::<String>(0)?, row.get::<String>(1)?);
479        }
480
481        // Load observation counts
482        let obs_count_sql = format!(
483            "SELECT entity_name, COUNT(*) FROM mcp_observations WHERE entity_name IN ({}) GROUP BY entity_name",
484            placeholders
485        );
486        let mut rows = conn.query(&obs_count_sql, params).await?;
487        while let Some(row) = rows.next().await? {
488            obs_counts.insert(row.get::<String>(0)?, row.get::<i64>(1)? as usize);
489        }
490    }
491
492    let mut entities = Vec::new();
493    for name in names {
494        if let Some(entity_type) = entity_types.get(name) {
495            let entity_truths = truths.get(name).cloned().unwrap_or_default();
496            let count = obs_counts.get(name).cloned().unwrap_or(0);
497            entities.push(crate::mcp::EntityOutput {
498                name: name.clone(),
499                entity_type: entity_type.clone(),
500                observations: Vec::new(),
501                truths: entity_truths,
502                observation_count: count,
503                body: None,
504            });
505        }
506    }
507    Ok(entities)
508}
509
510async fn load_entities_eager(
511    conn: &Connection,
512    names: &[String],
513) -> Result<Vec<crate::mcp::EntityOutput>> {
514    if names.is_empty() {
515        return Ok(Vec::new());
516    }
517
518    let mut entity_types = HashMap::new();
519    let mut observations: HashMap<String, Vec<String>> = HashMap::new();
520    let mut skill_bodies: HashMap<String, String> = HashMap::new();
521    let truths = select_truths(conn, names).await?;
522
523    for chunk in names.chunks(500) {
524        let placeholders = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(",");
525
526        // Load entities
527        let entity_sql =
528            crate::constant::SQL_SELECT_ENTITIES_IN_TEMPLATE.replace("{}", &placeholders);
529        let params = chunk
530            .iter()
531            .cloned()
532            .map(libsql::Value::from)
533            .collect::<Vec<_>>();
534
535        let mut rows = conn.query(&entity_sql, params.clone()).await?;
536        while let Some(row) = rows.next().await? {
537            entity_types.insert(row.get::<String>(0)?, row.get::<String>(1)?);
538        }
539
540        // Load observations
541        let obs_sql =
542            crate::constant::SQL_SELECT_OBSERVATIONS_IN_TEMPLATE.replace("{}", &placeholders);
543
544        let mut rows = conn.query(&obs_sql, params.clone()).await?;
545        while let Some(row) = rows.next().await? {
546            observations
547                .entry(row.get::<String>(0)?)
548                .or_default()
549                .push(row.get::<String>(1)?);
550        }
551
552        // Load skill bodies
553        let skill_sql =
554            crate::constant::SQL_SELECT_SKILL_BODIES_IN_TEMPLATE.replace("{}", &placeholders);
555        let mut rows = conn.query(&skill_sql, params).await?;
556        while let Some(row) = rows.next().await? {
557            skill_bodies.insert(row.get::<String>(0)?, row.get::<String>(1)?);
558        }
559    }
560
561    let mut entities = Vec::new();
562    for name in names {
563        if let Some(entity_type) = entity_types.get(name) {
564            let entity_truths = truths.get(name).cloned().unwrap_or_default();
565            let entity_obs = observations.remove(name).unwrap_or_default();
566            let count = entity_obs.len();
567            let body = skill_bodies.get(name).cloned();
568            entities.push(crate::mcp::EntityOutput {
569                name: name.clone(),
570                entity_type: entity_type.clone(),
571                observations: entity_obs,
572                truths: entity_truths,
573                observation_count: count,
574                body,
575            });
576        }
577    }
578    Ok(entities)
579}
580
581pub async fn mcp_stats(conn: &Connection) -> Result<(usize, usize, usize)> {
582    let mut rows = conn.query(crate::constant::SQL_COUNT_ENTITIES, ()).await?;
583    let entities_count: i64 = if let Some(row) = rows.next().await? {
584        row.get(0)?
585    } else {
586        0
587    };
588
589    let mut rows = conn.query(crate::constant::SQL_COUNT_RELATIONS, ()).await?;
590    let relations_count: i64 = if let Some(row) = rows.next().await? {
591        row.get(0)?
592    } else {
593        0
594    };
595
596    let mut rows = conn
597        .query(crate::constant::SQL_COUNT_OBSERVATIONS, ())
598        .await?;
599    let observations_count: i64 = if let Some(row) = rows.next().await? {
600        row.get(0)?
601    } else {
602        0
603    };
604
605    Ok((
606        entities_count as usize,
607        relations_count as usize,
608        observations_count as usize,
609    ))
610}
611
612pub async fn mcp_reset(conn: &Connection) -> Result<()> {
613    conn.execute(crate::constant::SQL_DELETE_ALL_RELATIONS, ())
614        .await?;
615    conn.execute(crate::constant::SQL_DELETE_ALL_OBSERVATIONS, ())
616        .await?;
617    conn.execute(crate::constant::SQL_DELETE_ALL_ENTITIES, ())
618        .await?;
619    Ok(())
620}
621
622pub async fn truth_upsert(conn: &Connection, entity: &str, key: &str, value: &str) -> Result<()> {
623    let norm_entity = crate::normalize::normalize_key(entity);
624    conn.execute(
625        crate::constant::SQL_UPSERT_TRUTH,
626        libsql::params![norm_entity, key, value],
627    )
628    .await?;
629    Ok(())
630}
631
632pub async fn truth_delete(conn: &Connection, entity: &str, key: &str) -> Result<()> {
633    let norm_entity = crate::normalize::normalize_key(entity);
634    conn.execute(
635        crate::constant::SQL_DELETE_TRUTH,
636        libsql::params![norm_entity, key],
637    )
638    .await?;
639    Ok(())
640}
641
642pub async fn select_truths(
643    conn: &Connection,
644    names: &[impl AsRef<str>],
645) -> Result<HashMap<String, std::collections::BTreeMap<String, String>>> {
646    let mut results = HashMap::new();
647    if names.is_empty() {
648        return Ok(results);
649    }
650
651    let normalized_names: Vec<String> = names
652        .iter()
653        .map(|n| crate::normalize::normalize_key(n.as_ref()))
654        .filter(|n| !n.is_empty())
655        .collect();
656
657    if normalized_names.is_empty() {
658        return Ok(results);
659    }
660
661    for chunk in normalized_names.chunks(500) {
662        let placeholders = chunk.iter().map(|_| "?").collect::<Vec<_>>().join(",");
663        let sql = crate::constant::SQL_SELECT_TRUTHS_FOR_ENTITIES.replace("{}", &placeholders);
664        let params = chunk
665            .iter()
666            .cloned()
667            .map(libsql::Value::from)
668            .collect::<Vec<_>>();
669
670        let mut rows = conn.query(&sql, params).await?;
671        while let Some(row) = rows.next().await? {
672            let entity_name: String = row.get(0)?;
673            let key: String = row.get(1)?;
674            let value: String = row.get(2)?;
675            results
676                .entry(entity_name)
677                .or_insert_with(std::collections::BTreeMap::new)
678                .insert(key, value);
679        }
680    }
681
682    Ok(results)
683}
684
685#[derive(Debug, Clone, PartialEq, Eq)]
686pub struct SkillRow {
687    pub entity_name: String,
688    pub description: String,
689    pub version: String,
690    pub source: String,
691    pub installed_at: String,
692}
693
694pub async fn skill_upsert(
695    conn: &Connection,
696    entity: &str,
697    body: &str,
698    source: &str,
699    version: &str,
700) -> Result<()> {
701    let norm_entity = crate::normalize::normalize_key(entity);
702    conn.execute(
703        crate::constant::SQL_UPSERT_SKILL,
704        libsql::params![norm_entity, body, source, version],
705    )
706    .await?;
707    Ok(())
708}
709
710pub async fn skill_body(conn: &Connection, entity: &str) -> Result<Option<String>> {
711    let norm_entity = crate::normalize::normalize_key(entity);
712    let mut rows = conn
713        .query(
714            crate::constant::SQL_SELECT_SKILL_BODY,
715            libsql::params![norm_entity],
716        )
717        .await?;
718    if let Some(row) = rows.next().await? {
719        let body: String = row.get(0)?;
720        Ok(Some(body))
721    } else {
722        Ok(None)
723    }
724}
725
726pub async fn list_skills(conn: &Connection) -> Result<Vec<SkillRow>> {
727    let mut rows = conn.query(crate::constant::SQL_LIST_SKILLS, ()).await?;
728    let mut results = Vec::new();
729    while let Some(row) = rows.next().await? {
730        results.push(SkillRow {
731            entity_name: row.get(0)?,
732            description: row.get(1)?,
733            version: row.get(2)?,
734            source: row.get(3)?,
735            installed_at: row.get(4)?,
736        });
737    }
738    Ok(results)
739}
740
741#[cfg(test)]
742mod tests {
743    use super::*;
744    use tempfile::tempdir;
745
746    #[tokio::test]
747    async fn test_init_creates_all_tables() {
748        let dir = tempdir().unwrap();
749        let db_path = dir.path().join("test.db");
750        unsafe {
751            std::env::set_var(ENV_DATABASE_URL, db_path.to_str().unwrap());
752        }
753        let (_db, conn) = init_db().await.unwrap();
754
755        // FTS5 table should exist
756        let mut rows = conn
757            .query("SELECT name FROM sqlite_master WHERE name='topics_fts'", ())
758            .await
759            .unwrap();
760        let row = rows.next().await.unwrap();
761        assert!(row.is_some(), "topics_fts table missing");
762    }
763
764    #[tokio::test]
765    async fn test_fts_search_finds_topic() {
766        let dir = tempdir().unwrap();
767        let db_path = dir.path().join("test.db");
768        unsafe {
769            std::env::set_var(ENV_DATABASE_URL, db_path.to_str().unwrap());
770        }
771        let (_db, conn) = init_db().await.unwrap();
772
773        conn.execute(
774            "INSERT INTO topics (id, title, file_path) VALUES ('rust-pin', 'Rust Pinning', '.miku/topics/rust-pinning.md')",
775            (),
776        ).await.unwrap();
777        conn.execute(
778            "INSERT INTO topics_fts (rowid, title, body) VALUES ((SELECT rowid FROM topics WHERE id='rust-pin'), 'Rust Pinning', 'pinning is a mechanism...')",
779            (),
780        ).await.unwrap();
781
782        let results = search_fts(&conn, "pinning", 5).await.unwrap();
783        assert_eq!(results.len(), 1);
784        assert_eq!(results[0].1, "Rust Pinning");
785    }
786
787    #[tokio::test]
788    async fn test_upsert_topic_preserves_created_at() {
789        let dir = tempdir().unwrap();
790        let db_path = dir.path().join("test.db");
791        unsafe {
792            std::env::set_var(ENV_DATABASE_URL, db_path.to_str().unwrap());
793        }
794        let (_db, conn) = init_db().await.unwrap();
795
796        // Seed a row with a fixed, distinguishable created_at so a reset would
797        // be detectable even within the same wall-clock second.
798        conn.execute(
799            "INSERT INTO topics (id, title, file_path, body, created_at) \
800             VALUES ('t1', 'Old Title', '/old', 'old body', '2000-01-01 00:00:00')",
801            (),
802        )
803        .await
804        .unwrap();
805
806        // Re-upsert the same id with new content.
807        upsert_topic(&conn, "t1", "New Title", "/new", "new body")
808            .await
809            .unwrap();
810
811        let mut rows = conn
812            .query(
813                "SELECT title, body, created_at FROM topics WHERE id = 't1'",
814                (),
815            )
816            .await
817            .unwrap();
818        let row = rows.next().await.unwrap().expect("row should exist");
819        let title: String = row.get(0).unwrap();
820        let body: String = row.get(1).unwrap();
821        let created_at: String = row.get(2).unwrap();
822
823        // Update was applied...
824        assert_eq!(title, "New Title");
825        assert_eq!(body, "new body");
826        // ...but created_at must be preserved (INSERT OR REPLACE would reset it).
827        assert_eq!(created_at, "2000-01-01 00:00:00");
828    }
829
830    async fn seed_entity(conn: &Connection, name: &str, entity_type: &str, obs: &[&str]) {
831        mcp_create_entities(
832            conn,
833            vec![crate::mcp::EntityInput {
834                name: name.to_string(),
835                entity_type: entity_type.to_string(),
836                observations: obs.iter().map(|s| s.to_string()).collect(),
837            }],
838        )
839        .await
840        .unwrap();
841    }
842
843    #[tokio::test]
844    async fn test_search_nodes_stemming() {
845        let dir = tempdir().unwrap();
846        unsafe {
847            std::env::set_var(
848                ENV_DATABASE_URL,
849                dir.path().join("test.db").to_str().unwrap(),
850            );
851        }
852        let (_db, conn) = init_db().await.unwrap();
853
854        seed_entity(
855            &conn,
856            "async-patterns",
857            "concept",
858            &["running async tasks efficiently"],
859        )
860        .await;
861
862        // "run" should match "running" via porter stemming
863        let graph = mcp_search_nodes(&conn, "run").await.unwrap();
864        assert_eq!(graph.entities.len(), 1);
865        assert_eq!(graph.entities[0].name, "async-patterns");
866    }
867
868    #[tokio::test]
869    async fn test_search_nodes_entity_name_fallback() {
870        let dir = tempdir().unwrap();
871        unsafe {
872            std::env::set_var(
873                ENV_DATABASE_URL,
874                dir.path().join("test.db").to_str().unwrap(),
875            );
876        }
877        let (_db, conn) = init_db().await.unwrap();
878
879        // Entity with no observations — FTS finds nothing, LIKE fallback finds by name
880        seed_entity(&conn, "user-preferences", "preference", &[]).await;
881
882        let graph = mcp_search_nodes(&conn, "user-preferences").await.unwrap();
883        assert_eq!(graph.entities.len(), 1);
884        assert_eq!(graph.entities[0].name, "user-preferences");
885    }
886
887    #[tokio::test]
888    async fn test_search_nodes_bm25_ordering() {
889        let dir = tempdir().unwrap();
890        unsafe {
891            std::env::set_var(
892                ENV_DATABASE_URL,
893                dir.path().join("test.db").to_str().unwrap(),
894            );
895        }
896        let (_db, conn) = init_db().await.unwrap();
897
898        // "alpha" has both query words; "beta" has only one — alpha should rank first
899        seed_entity(&conn, "alpha", "project", &["async tokio runtime patterns"]).await;
900        seed_entity(&conn, "beta", "project", &["tokio scheduler"]).await;
901
902        let graph = mcp_search_nodes(&conn, "async tokio").await.unwrap();
903        assert!(!graph.entities.is_empty());
904        assert_eq!(graph.entities[0].name, "alpha");
905    }
906
907    #[tokio::test]
908    async fn test_search_nodes_invalid_fts_syntax_no_panic() {
909        let dir = tempdir().unwrap();
910        unsafe {
911            std::env::set_var(
912                ENV_DATABASE_URL,
913                dir.path().join("test.db").to_str().unwrap(),
914            );
915        }
916        let (_db, conn) = init_db().await.unwrap();
917
918        // Invalid FTS5 syntax — must not panic, falls back to LIKE gracefully
919        let result = mcp_search_nodes(&conn, "AND AND").await;
920        assert!(result.is_ok());
921    }
922
923    #[tokio::test]
924    async fn test_search_nodes_default_limit_and_explicit_limit() {
925        let dir = tempdir().unwrap();
926        unsafe {
927            std::env::set_var(
928                ENV_DATABASE_URL,
929                dir.path().join("test.db").to_str().unwrap(),
930            );
931        }
932        let (_db, conn) = init_db().await.unwrap();
933
934        for i in 0..(DEFAULT_SEARCH_LIMIT + 10) {
935            seed_entity(&conn, &format!("entity-{i:03}"), "project", &["commonterm"]).await;
936        }
937
938        let default_graph = mcp_search_nodes(&conn, "commonterm").await.unwrap();
939        assert_eq!(default_graph.entities.len(), DEFAULT_SEARCH_LIMIT);
940
941        let explicit_graph = mcp_search_nodes_with_limit(&conn, "commonterm", 7)
942            .await
943            .unwrap();
944        assert_eq!(explicit_graph.entities.len(), 7);
945    }
946
947    #[tokio::test]
948    async fn test_mcp_stats_and_reset() {
949        let dir = tempdir().unwrap();
950        unsafe {
951            std::env::set_var(
952                ENV_DATABASE_URL,
953                dir.path().join("test.db").to_str().unwrap(),
954            );
955        }
956        let (_db, conn) = init_db().await.unwrap();
957
958        // Check empty stats
959        let stats = mcp_stats(&conn).await.unwrap();
960        assert_eq!(stats, (0, 0, 0));
961
962        // Seed some data
963        seed_entity(&conn, "entity1", "project", &["obs1", "obs2"]).await;
964        seed_entity(&conn, "entity2", "project", &["obs3"]).await;
965        mcp_create_relations(
966            &conn,
967            vec![crate::mcp::RelationInput {
968                from: "entity1".to_string(),
969                to: "entity2".to_string(),
970                relation_type: "related".to_string(),
971            }],
972        )
973        .await
974        .unwrap();
975
976        // Check populated stats
977        let stats = mcp_stats(&conn).await.unwrap();
978        assert_eq!(stats, (2, 1, 3));
979
980        // Test reset
981        mcp_reset(&conn).await.unwrap();
982
983        // Check empty stats again
984        let stats = mcp_stats(&conn).await.unwrap();
985        assert_eq!(stats, (0, 0, 0));
986    }
987
988    #[tokio::test]
989    async fn test_truth_upsert_twice_same_key() {
990        let dir = tempdir().unwrap();
991        unsafe {
992            std::env::set_var(
993                ENV_DATABASE_URL,
994                dir.path().join("test.db").to_str().unwrap(),
995            );
996        }
997        let (_db, conn) = init_db().await.unwrap();
998
999        seed_entity(&conn, "test-entity", "concept", &[]).await;
1000
1001        truth_upsert(&conn, "test-entity", "version", "1.0.0")
1002            .await
1003            .unwrap();
1004        truth_upsert(&conn, "test-entity", "version", "1.0.1")
1005            .await
1006            .unwrap();
1007
1008        let truths = select_truths(&conn, &["test-entity"]).await.unwrap();
1009        let entity_truths = truths.get("test-entity").expect("should have truths");
1010        assert_eq!(entity_truths.len(), 1);
1011        assert_eq!(entity_truths.get("version").unwrap(), "1.0.1");
1012    }
1013
1014    #[tokio::test]
1015    async fn test_truth_upsert_two_keys() {
1016        let dir = tempdir().unwrap();
1017        unsafe {
1018            std::env::set_var(
1019                ENV_DATABASE_URL,
1020                dir.path().join("test.db").to_str().unwrap(),
1021            );
1022        }
1023        let (_db, conn) = init_db().await.unwrap();
1024
1025        seed_entity(&conn, "test-entity", "concept", &[]).await;
1026
1027        truth_upsert(&conn, "test-entity", "version", "1.0.0")
1028            .await
1029            .unwrap();
1030        truth_upsert(&conn, "test-entity", "author", "Alice")
1031            .await
1032            .unwrap();
1033
1034        let truths = select_truths(&conn, &["test-entity"]).await.unwrap();
1035        let entity_truths = truths.get("test-entity").expect("should have truths");
1036        assert_eq!(entity_truths.len(), 2);
1037        assert_eq!(entity_truths.get("author").unwrap(), "Alice");
1038        assert_eq!(entity_truths.get("version").unwrap(), "1.0.0");
1039    }
1040
1041    #[tokio::test]
1042    async fn test_truth_delete() {
1043        let dir = tempdir().unwrap();
1044        unsafe {
1045            std::env::set_var(
1046                ENV_DATABASE_URL,
1047                dir.path().join("test.db").to_str().unwrap(),
1048            );
1049        }
1050        let (_db, conn) = init_db().await.unwrap();
1051
1052        seed_entity(&conn, "test-entity", "concept", &[]).await;
1053
1054        truth_upsert(&conn, "test-entity", "k1", "v1")
1055            .await
1056            .unwrap();
1057        truth_upsert(&conn, "test-entity", "k2", "v2")
1058            .await
1059            .unwrap();
1060
1061        truth_delete(&conn, "test-entity", "k1").await.unwrap();
1062
1063        let truths = select_truths(&conn, &["test-entity"]).await.unwrap();
1064        let entity_truths = truths.get("test-entity").expect("should have truths");
1065        assert_eq!(entity_truths.len(), 1);
1066        assert_eq!(entity_truths.get("k2").unwrap(), "v2");
1067    }
1068
1069    #[tokio::test]
1070    async fn test_delete_entities_cascades_truths() {
1071        let dir = tempdir().unwrap();
1072        unsafe {
1073            std::env::set_var(
1074                ENV_DATABASE_URL,
1075                dir.path().join("test.db").to_str().unwrap(),
1076            );
1077        }
1078        let (_db, conn) = init_db().await.unwrap();
1079
1080        seed_entity(&conn, "test-entity", "concept", &[]).await;
1081        truth_upsert(&conn, "test-entity", "k1", "v1")
1082            .await
1083            .unwrap();
1084
1085        mcp_delete_entities(&conn, vec!["test-entity".to_string()])
1086            .await
1087            .unwrap();
1088
1089        // Check if the truth was deleted.
1090        let mut rows = conn
1091            .query("SELECT COUNT(*) FROM mcp_truths", ())
1092            .await
1093            .unwrap();
1094        let count: i64 = rows.next().await.unwrap().unwrap().get(0).unwrap();
1095        assert_eq!(count, 0);
1096    }
1097
1098    #[tokio::test]
1099    async fn test_observation_limit_evicts_oldest() {
1100        let dir = tempdir().unwrap();
1101        unsafe {
1102            std::env::set_var(
1103                ENV_DATABASE_URL,
1104                dir.path().join("test.db").to_str().unwrap(),
1105            );
1106        }
1107        let (_db, conn) = init_db().await.unwrap();
1108
1109        seed_entity(&conn, "test-entity", "concept", &[]).await;
1110
1111        let inputs = vec![crate::mcp::ObservationInput {
1112            entity_name: "test-entity".to_string(),
1113            contents: vec![
1114                "obs1".to_string(),
1115                "obs2".to_string(),
1116                "obs3".to_string(),
1117                "obs4".to_string(),
1118                "obs5".to_string(),
1119            ],
1120        }];
1121        mcp_add_observations(&conn, inputs, 3).await.unwrap();
1122
1123        let graph = mcp_open_nodes(&conn, vec!["test-entity".to_string()])
1124            .await
1125            .unwrap();
1126        let entity = &graph.entities[0];
1127        let mut obs = entity.observations.clone();
1128        obs.sort();
1129        assert_eq!(
1130            obs,
1131            vec!["obs3".to_string(), "obs4".to_string(), "obs5".to_string(),]
1132        );
1133    }
1134
1135    #[tokio::test]
1136    async fn test_observation_limit_zero_is_unbounded() {
1137        let dir = tempdir().unwrap();
1138        unsafe {
1139            std::env::set_var(
1140                ENV_DATABASE_URL,
1141                dir.path().join("test.db").to_str().unwrap(),
1142            );
1143        }
1144        let (_db, conn) = init_db().await.unwrap();
1145
1146        seed_entity(&conn, "test-entity", "concept", &[]).await;
1147
1148        let inputs = vec![crate::mcp::ObservationInput {
1149            entity_name: "test-entity".to_string(),
1150            contents: vec![
1151                "obs1".to_string(),
1152                "obs2".to_string(),
1153                "obs3".to_string(),
1154                "obs4".to_string(),
1155                "obs5".to_string(),
1156            ],
1157        }];
1158        mcp_add_observations(&conn, inputs, 0).await.unwrap();
1159
1160        let graph = mcp_open_nodes(&conn, vec!["test-entity".to_string()])
1161            .await
1162            .unwrap();
1163        let entity = &graph.entities[0];
1164        assert_eq!(entity.observations.len(), 5);
1165    }
1166
1167    #[tokio::test]
1168    async fn test_lazy_read_graph_and_search_nodes() {
1169        let dir = tempdir().unwrap();
1170        unsafe {
1171            std::env::set_var(
1172                ENV_DATABASE_URL,
1173                dir.path().join("test.db").to_str().unwrap(),
1174            );
1175        }
1176        let (_db, conn) = init_db().await.unwrap();
1177
1178        seed_entity(&conn, "test-entity", "concept", &["obs1", "obs2"]).await;
1179        truth_upsert(&conn, "test-entity", "k1", "v1")
1180            .await
1181            .unwrap();
1182
1183        // 1. Test read-graph (should be lazy)
1184        let graph_read = mcp_read_graph(&conn).await.unwrap();
1185        let entity_read = &graph_read.entities[0];
1186        assert!(entity_read.observations.is_empty());
1187        assert_eq!(entity_read.observation_count, 2);
1188        assert_eq!(entity_read.truths.len(), 1);
1189        assert_eq!(entity_read.truths.get("k1").unwrap(), "v1");
1190
1191        // 2. Test search-nodes (should be lazy)
1192        let graph_search = mcp_search_nodes(&conn, "test").await.unwrap();
1193        let entity_search = &graph_search.entities[0];
1194        assert!(entity_search.observations.is_empty());
1195        assert_eq!(entity_search.observation_count, 2);
1196        assert_eq!(entity_search.truths.len(), 1);
1197        assert_eq!(entity_search.truths.get("k1").unwrap(), "v1");
1198
1199        // 3. Test open-nodes (should be eager)
1200        let graph_open = mcp_open_nodes(&conn, vec!["test-entity".to_string()])
1201            .await
1202            .unwrap();
1203        let entity_open = &graph_open.entities[0];
1204        assert_eq!(entity_open.observations.len(), 2);
1205        assert_eq!(entity_open.observation_count, 2);
1206        assert_eq!(entity_open.truths.len(), 1);
1207        assert_eq!(entity_open.truths.get("k1").unwrap(), "v1");
1208    }
1209
1210    #[tokio::test]
1211    async fn test_skill_storage_and_mcp_open_nodes() {
1212        let dir = tempdir().unwrap();
1213        unsafe {
1214            std::env::set_var(
1215                ENV_DATABASE_URL,
1216                dir.path().join("test.db").to_str().unwrap(),
1217            );
1218        }
1219        let (_db, conn) = init_db().await.unwrap();
1220
1221        seed_entity(&conn, "skill:test-skill", "skill", &[]).await;
1222        truth_upsert(&conn, "skill:test-skill", "description", "my test skill")
1223            .await
1224            .unwrap();
1225
1226        // 1. Upsert skill
1227        skill_upsert(
1228            &conn,
1229            "skill:test-skill",
1230            "body content 1",
1231            "source-1",
1232            "1.0.0",
1233        )
1234        .await
1235        .unwrap();
1236
1237        // 2. open-nodes should return the body
1238        let graph = mcp_open_nodes(&conn, vec!["skill:test-skill".to_string()])
1239            .await
1240            .unwrap();
1241        let entity = &graph.entities[0];
1242        assert_eq!(entity.body.as_deref(), Some("body content 1"));
1243
1244        // 3. read-graph and search-nodes should NOT return the body
1245        let graph_read = mcp_read_graph(&conn).await.unwrap();
1246        assert!(graph_read.entities[0].body.is_none());
1247
1248        let graph_search = mcp_search_nodes(&conn, "skill").await.unwrap();
1249        assert!(graph_search.entities[0].body.is_none());
1250
1251        // 4. Second upsert should replace the body
1252        skill_upsert(
1253            &conn,
1254            "skill:test-skill",
1255            "body content 2",
1256            "source-1",
1257            "1.0.1",
1258        )
1259        .await
1260        .unwrap();
1261        let graph_2 = mcp_open_nodes(&conn, vec!["skill:test-skill".to_string()])
1262            .await
1263            .unwrap();
1264        assert_eq!(graph_2.entities[0].body.as_deref(), Some("body content 2"));
1265
1266        // 5. list_skills should list name + description + version + source + installed_at
1267        let skills = list_skills(&conn).await.unwrap();
1268        assert_eq!(skills.len(), 1);
1269        assert_eq!(skills[0].entity_name, "skill:test-skill");
1270        assert_eq!(skills[0].description, "my test skill");
1271        assert_eq!(skills[0].version, "1.0.1");
1272        assert_eq!(skills[0].source, "source-1");
1273
1274        // 6. delete-entities cascades skills
1275        mcp_delete_entities(&conn, vec!["skill:test-skill".to_string()])
1276            .await
1277            .unwrap();
1278        let body_after = skill_body(&conn, "skill:test-skill").await.unwrap();
1279        assert!(body_after.is_none());
1280
1281        let count_skills: i64 = conn
1282            .query("SELECT COUNT(*) FROM mcp_skills", ())
1283            .await
1284            .unwrap()
1285            .next()
1286            .await
1287            .unwrap()
1288            .unwrap()
1289            .get(0)
1290            .unwrap();
1291        assert_eq!(count_skills, 0);
1292    }
1293}
miku/db.rs

miku/
db.rs
