context_harness/
app_store.rs

1//! Application-level storage boundary for native Context Harness operations.
2//!
3//! [`context_harness_core::store::Store`] covers reusable search and retrieval
4//! behavior. [`AppStore`] adds CLI/native responsibilities that are still
5//! canonical SQLite operations today: migrations, sync checkpoints, source item
6//! writes, embedding maintenance, stats, and export views.
7
8use anyhow::Result;
9use async_trait::async_trait;
10use serde::Serialize;
11use sha2::{Digest, Sha256};
12use sqlx::{Row, SqlitePool};
13use uuid::Uuid;
14
15use context_harness_core::models::{Chunk, Document};
16use context_harness_core::store::{ChunkCandidate, DocumentMetadata, DocumentResponse, Store};
17
18use crate::config::Config;
19use crate::db;
20use crate::migrate;
21use crate::models::SourceItem;
22use crate::sqlite_store::SqliteStore;
23use crate::vector_index::{self, VectorRecord};
24
25/// A chunk that needs embedding because its embedding is missing or stale.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct PendingChunk {
28    pub chunk_id: String,
29    pub document_id: String,
30    pub text: String,
31    pub text_hash: String,
32}
33
34/// Per-source database statistics.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct SourceStats {
37    pub source: String,
38    pub doc_count: i64,
39    pub chunk_count: i64,
40    pub embedded_count: i64,
41    pub last_sync_ts: Option<i64>,
42}
43
44/// Database statistics used by `ctx stats`.
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct StoreStats {
47    pub total_docs: i64,
48    pub total_chunks: i64,
49    pub total_embedded: i64,
50    pub db_size_bytes: u64,
51    pub sources: Vec<SourceStats>,
52}
53
54/// Export payload used by `ctx export`.
55#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
56pub struct ExportData {
57    pub documents: Vec<ExportDocument>,
58    pub chunks: Vec<ExportChunk>,
59}
60
61/// Exported document row.
62#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
63pub struct ExportDocument {
64    pub id: String,
65    pub source: String,
66    pub source_id: String,
67    pub source_url: Option<String>,
68    pub title: Option<String>,
69    pub updated_at: i64,
70    pub body: String,
71}
72
73/// Exported chunk row.
74#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
75pub struct ExportChunk {
76    pub id: String,
77    pub document_id: String,
78    pub chunk_index: i64,
79    pub text: String,
80}
81
82/// App-level storage operations layered on top of core search storage.
83#[async_trait]
84pub trait AppStore: Store {
85    #[allow(dead_code)]
86    async fn initialize(&self) -> Result<()>;
87    async fn get_checkpoint(&self, source: &str) -> Result<Option<i64>>;
88    async fn set_checkpoint(&self, source: &str, cursor: i64) -> Result<()>;
89    async fn upsert_source_item(&self, item: &SourceItem) -> Result<String>;
90    async fn find_pending_chunks(
91        &self,
92        model: &str,
93        limit: Option<usize>,
94    ) -> Result<Vec<PendingChunk>>;
95    async fn get_embedding_hash(&self, chunk_id: &str, model: &str) -> Result<Option<String>>;
96    async fn clear_embeddings(&self) -> Result<()>;
97    async fn stats(&self) -> Result<StoreStats>;
98    async fn export_index(&self) -> Result<ExportData>;
99}
100
101/// SQLite-backed [`AppStore`] implementation.
102pub struct SqliteAppStore {
103    config: Config,
104    pool: SqlitePool,
105}
106
107impl SqliteAppStore {
108    pub async fn connect(config: &Config) -> Result<Self> {
109        let pool = db::connect(config).await?;
110        Ok(Self {
111            config: config.clone(),
112            pool,
113        })
114    }
115
116    pub async fn initialize_config(config: &Config) -> Result<()> {
117        migrate::run_migrations(config).await
118    }
119
120    #[allow(dead_code)]
121    pub fn pool(&self) -> &SqlitePool {
122        &self.pool
123    }
124
125    /// Runtime history uses the same canonical database as ingestion/retrieval.
126    #[allow(dead_code)]
127    pub fn agent_runs(&self, workspace_id: &str) -> Result<crate::agent_store::AgentRunStore> {
128        crate::agent_store::AgentRunStore::new(self.pool.clone(), workspace_id)
129    }
130
131    /// Accepted background tasks share the canonical database but remain
132    /// workspace-scoped and distinct from concrete run attempts.
133    pub fn agent_tasks(
134        &self,
135        workspace_id: &str,
136    ) -> Result<crate::agent_task_store::AgentTaskStore> {
137        crate::agent_task_store::AgentTaskStore::new(self.pool.clone(), workspace_id)
138    }
139
140    pub async fn close(&self) {
141        self.pool.close().await;
142    }
143
144    fn core_store(&self) -> SqliteStore {
145        SqliteStore::new(self.pool.clone())
146    }
147
148    /// Atomically replace one canonical document and all of its searchable chunks.
149    #[allow(dead_code)] // The binary target compiles this module without the public library API.
150    pub(crate) async fn replace_source_item_atomic(
151        &self,
152        item: &SourceItem,
153        max_tokens: usize,
154    ) -> Result<CanonicalDocumentWrite> {
155        let mut tx = self.pool.begin().await?;
156        let existing_id: Option<String> =
157            sqlx::query_scalar("SELECT id FROM documents WHERE source = ? AND source_id = ?")
158                .bind(&item.source)
159                .bind(&item.source_id)
160                .fetch_optional(&mut *tx)
161                .await?;
162        let document = document_from_source_item(
163            existing_id.unwrap_or_else(|| Uuid::new_v4().to_string()),
164            item,
165        );
166        let chunks = crate::chunk::chunk_text(&document.id, &document.body, max_tokens);
167
168        sqlx::query(
169            r#"
170            INSERT INTO documents (id, source, source_id, source_url, title, author,
171                                   created_at, updated_at, content_type, body,
172                                   metadata_json, raw_json, dedup_hash)
173            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
174            ON CONFLICT(source, source_id) DO UPDATE SET
175                source_url = excluded.source_url,
176                title = excluded.title,
177                author = excluded.author,
178                updated_at = excluded.updated_at,
179                content_type = excluded.content_type,
180                body = excluded.body,
181                metadata_json = excluded.metadata_json,
182                raw_json = excluded.raw_json,
183                dedup_hash = excluded.dedup_hash
184            "#,
185        )
186        .bind(&document.id)
187        .bind(&document.source)
188        .bind(&document.source_id)
189        .bind(&document.source_url)
190        .bind(&document.title)
191        .bind(&document.author)
192        .bind(document.created_at)
193        .bind(document.updated_at)
194        .bind(&document.content_type)
195        .bind(&document.body)
196        .bind(&document.metadata_json)
197        .bind(&document.raw_json)
198        .bind(&document.dedup_hash)
199        .execute(&mut *tx)
200        .await?;
201
202        sqlx::query(
203            "DELETE FROM chunk_vectors WHERE chunk_id IN (SELECT id FROM chunks WHERE document_id = ?)",
204        )
205        .bind(&document.id)
206        .execute(&mut *tx)
207        .await?;
208        sqlx::query(
209            "DELETE FROM embeddings WHERE chunk_id IN (SELECT id FROM chunks WHERE document_id = ?)",
210        )
211        .bind(&document.id)
212        .execute(&mut *tx)
213        .await?;
214        sqlx::query("DELETE FROM chunks_fts WHERE document_id = ?")
215            .bind(&document.id)
216            .execute(&mut *tx)
217            .await?;
218        sqlx::query("DELETE FROM chunks WHERE document_id = ?")
219            .bind(&document.id)
220            .execute(&mut *tx)
221            .await?;
222        for chunk in &chunks {
223            sqlx::query(
224                "INSERT INTO chunks (id, document_id, chunk_index, text, hash) VALUES (?, ?, ?, ?, ?)",
225            )
226            .bind(&chunk.id)
227            .bind(&chunk.document_id)
228            .bind(chunk.chunk_index)
229            .bind(&chunk.text)
230            .bind(&chunk.hash)
231            .execute(&mut *tx)
232            .await?;
233            sqlx::query("INSERT INTO chunks_fts (chunk_id, document_id, text) VALUES (?, ?, ?)")
234                .bind(&chunk.id)
235                .bind(&chunk.document_id)
236                .bind(&chunk.text)
237                .execute(&mut *tx)
238                .await?;
239        }
240        tx.commit().await?;
241
242        let sidecar_current = vector_index::remove_configured_sidecar(&self.config).is_ok();
243        Ok(CanonicalDocumentWrite {
244            document_id: document.id,
245            chunks,
246            sidecar_current,
247        })
248    }
249}
250
251#[allow(dead_code)] // The binary target compiles this module without the public library API.
252pub(crate) struct CanonicalDocumentWrite {
253    pub(crate) document_id: String,
254    pub(crate) chunks: Vec<Chunk>,
255    pub(crate) sidecar_current: bool,
256}
257
258#[async_trait]
259impl Store for SqliteAppStore {
260    async fn upsert_document(&self, doc: &Document) -> Result<String> {
261        self.core_store().upsert_document(doc).await
262    }
263
264    async fn replace_chunks(
265        &self,
266        doc_id: &str,
267        chunks: &[Chunk],
268        vectors: Option<&[Vec<f32>]>,
269    ) -> Result<()> {
270        self.core_store()
271            .replace_chunks(doc_id, chunks, vectors)
272            .await?;
273        vector_index::remove_configured_sidecar(&self.config)?;
274        Ok(())
275    }
276
277    async fn upsert_embedding(
278        &self,
279        chunk_id: &str,
280        doc_id: &str,
281        vector: &[f32],
282        model: &str,
283        dims: usize,
284        content_hash: &str,
285    ) -> Result<()> {
286        self.core_store()
287            .upsert_embedding(chunk_id, doc_id, vector, model, dims, content_hash)
288            .await?;
289        let record = VectorRecord {
290            chunk_id: chunk_id.to_string(),
291            document_id: doc_id.to_string(),
292            vector: vector.to_vec(),
293            model: model.to_string(),
294            dims,
295            content_hash: content_hash.to_string(),
296        };
297        vector_index::sync_vector_record_after_sqlite(&self.config, &self.pool, &record).await?;
298        Ok(())
299    }
300
301    async fn get_document(&self, id: &str) -> Result<Option<DocumentResponse>> {
302        self.core_store().get_document(id).await
303    }
304
305    async fn get_document_metadata(&self, id: &str) -> Result<Option<DocumentMetadata>> {
306        self.core_store().get_document_metadata(id).await
307    }
308
309    async fn keyword_search(
310        &self,
311        query: &str,
312        limit: i64,
313        source: Option<&str>,
314        since: Option<&str>,
315    ) -> Result<Vec<ChunkCandidate>> {
316        self.core_store()
317            .keyword_search(query, limit, source, since)
318            .await
319    }
320
321    async fn vector_search(
322        &self,
323        query_vec: &[f32],
324        limit: i64,
325        source: Option<&str>,
326        since: Option<&str>,
327    ) -> Result<Vec<ChunkCandidate>> {
328        self.core_store()
329            .vector_search(query_vec, limit, source, since)
330            .await
331    }
332}
333
334#[async_trait]
335impl AppStore for SqliteAppStore {
336    async fn initialize(&self) -> Result<()> {
337        migrate::run_migrations(&self.config).await
338    }
339
340    async fn get_checkpoint(&self, source: &str) -> Result<Option<i64>> {
341        let result: Option<String> =
342            sqlx::query_scalar("SELECT cursor FROM checkpoints WHERE source = ?")
343                .bind(source)
344                .fetch_optional(&self.pool)
345                .await?;
346
347        Ok(result.and_then(|s| s.parse::<i64>().ok()))
348    }
349
350    async fn set_checkpoint(&self, source: &str, cursor: i64) -> Result<()> {
351        let now = chrono::Utc::now().timestamp();
352        sqlx::query(
353            r#"
354            INSERT INTO checkpoints (source, cursor, updated_at) VALUES (?, ?, ?)
355            ON CONFLICT(source) DO UPDATE SET cursor = excluded.cursor, updated_at = excluded.updated_at
356            "#,
357        )
358        .bind(source)
359        .bind(cursor.to_string())
360        .bind(now)
361        .execute(&self.pool)
362        .await?;
363
364        Ok(())
365    }
366
367    async fn upsert_source_item(&self, item: &SourceItem) -> Result<String> {
368        let doc = source_item_to_document(&self.pool, item).await?;
369        self.upsert_document(&doc).await
370    }
371
372    async fn find_pending_chunks(
373        &self,
374        model: &str,
375        limit: Option<usize>,
376    ) -> Result<Vec<PendingChunk>> {
377        let limit_val = limit.unwrap_or(usize::MAX) as i64;
378
379        let rows = sqlx::query(
380            r#"
381            SELECT c.id AS chunk_id, c.document_id, c.text, c.hash AS chunk_hash
382            FROM chunks c
383            LEFT JOIN embeddings e ON e.chunk_id = c.id AND e.model = ?
384            WHERE e.chunk_id IS NULL OR e.hash != c.hash
385            ORDER BY c.document_id, c.chunk_index
386            LIMIT ?
387            "#,
388        )
389        .bind(model)
390        .bind(limit_val)
391        .fetch_all(&self.pool)
392        .await?;
393
394        Ok(rows
395            .iter()
396            .map(|row| {
397                let text: String = row.get("text");
398                PendingChunk {
399                    chunk_id: row.get("chunk_id"),
400                    document_id: row.get("document_id"),
401                    text_hash: row.get("chunk_hash"),
402                    text,
403                }
404            })
405            .collect())
406    }
407
408    async fn get_embedding_hash(&self, chunk_id: &str, model: &str) -> Result<Option<String>> {
409        let hash =
410            sqlx::query_scalar("SELECT hash FROM embeddings WHERE chunk_id = ? AND model = ?")
411                .bind(chunk_id)
412                .bind(model)
413                .fetch_optional(&self.pool)
414                .await?;
415        Ok(hash)
416    }
417
418    async fn clear_embeddings(&self) -> Result<()> {
419        sqlx::query("DELETE FROM chunk_vectors")
420            .execute(&self.pool)
421            .await?;
422        sqlx::query("DELETE FROM embeddings")
423            .execute(&self.pool)
424            .await?;
425        vector_index::remove_configured_sidecar(&self.config)?;
426        Ok(())
427    }
428
429    async fn stats(&self) -> Result<StoreStats> {
430        let total_docs: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM documents")
431            .fetch_one(&self.pool)
432            .await?;
433        let total_chunks: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM chunks")
434            .fetch_one(&self.pool)
435            .await?;
436        let total_embedded: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM chunk_vectors")
437            .fetch_one(&self.pool)
438            .await?;
439        let db_size_bytes = std::fs::metadata(&self.config.db.path)
440            .map(|m| m.len())
441            .unwrap_or(0);
442
443        let source_rows = sqlx::query(
444            r#"
445            SELECT
446                d.source,
447                COUNT(DISTINCT d.id) AS doc_count,
448                COUNT(DISTINCT c.id) AS chunk_count,
449                COUNT(DISTINCT cv.chunk_id) AS embedded_count
450            FROM documents d
451            LEFT JOIN chunks c ON c.document_id = d.id
452            LEFT JOIN chunk_vectors cv ON cv.chunk_id = c.id
453            GROUP BY d.source
454            ORDER BY doc_count DESC
455            "#,
456        )
457        .fetch_all(&self.pool)
458        .await?;
459
460        let checkpoint_rows = sqlx::query("SELECT source, updated_at FROM checkpoints")
461            .fetch_all(&self.pool)
462            .await?;
463
464        let mut sources = Vec::new();
465        for row in &source_rows {
466            let source: String = row.get("source");
467            let last_sync_ts = checkpoint_rows
468                .iter()
469                .find(|cp| {
470                    let cp_source: String = cp.get("source");
471                    cp_source == source
472                })
473                .map(|cp| cp.get::<i64, _>("updated_at"));
474
475            sources.push(SourceStats {
476                source,
477                doc_count: row.get("doc_count"),
478                chunk_count: row.get("chunk_count"),
479                embedded_count: row.get("embedded_count"),
480                last_sync_ts,
481            });
482        }
483
484        Ok(StoreStats {
485            total_docs,
486            total_chunks,
487            total_embedded,
488            db_size_bytes,
489            sources,
490        })
491    }
492
493    async fn export_index(&self) -> Result<ExportData> {
494        let doc_rows = sqlx::query(
495            "SELECT id, source, source_id, source_url, title, updated_at, body \
496             FROM documents ORDER BY source_id",
497        )
498        .fetch_all(&self.pool)
499        .await?;
500
501        let chunk_rows = sqlx::query(
502            "SELECT id, document_id, chunk_index, text \
503             FROM chunks ORDER BY document_id, chunk_index",
504        )
505        .fetch_all(&self.pool)
506        .await?;
507
508        let documents = doc_rows
509            .iter()
510            .map(|row| ExportDocument {
511                id: row.get("id"),
512                source: row.get("source"),
513                source_id: row.get("source_id"),
514                source_url: row.get("source_url"),
515                title: row.get("title"),
516                updated_at: row.get("updated_at"),
517                body: row.get("body"),
518            })
519            .collect();
520
521        let chunks = chunk_rows
522            .iter()
523            .map(|row| ExportChunk {
524                id: row.get("id"),
525                document_id: row.get("document_id"),
526                chunk_index: row.get("chunk_index"),
527                text: row.get("text"),
528            })
529            .collect();
530
531        Ok(ExportData { documents, chunks })
532    }
533}
534
535async fn source_item_to_document(pool: &SqlitePool, item: &SourceItem) -> Result<Document> {
536    let existing_id: Option<String> =
537        sqlx::query_scalar("SELECT id FROM documents WHERE source = ? AND source_id = ?")
538            .bind(&item.source)
539            .bind(&item.source_id)
540            .fetch_optional(pool)
541            .await?;
542
543    Ok(document_from_source_item(
544        existing_id.unwrap_or_else(|| Uuid::new_v4().to_string()),
545        item,
546    ))
547}
548
549fn document_from_source_item(id: String, item: &SourceItem) -> Document {
550    Document {
551        id,
552        source: item.source.clone(),
553        source_id: item.source_id.clone(),
554        source_url: item.source_url.clone(),
555        title: item.title.clone(),
556        author: item.author.clone(),
557        created_at: item.created_at.timestamp(),
558        updated_at: item.updated_at.timestamp(),
559        content_type: item.content_type.clone(),
560        body: item.body.clone(),
561        metadata_json: item.metadata_json.clone(),
562        raw_json: item.raw_json.clone(),
563        dedup_hash: dedup_hash(item),
564    }
565}
566
567fn dedup_hash(item: &SourceItem) -> String {
568    let mut hasher = Sha256::new();
569    hasher.update(item.source.as_bytes());
570    hasher.update(item.source_id.as_bytes());
571    hasher.update(item.updated_at.timestamp().to_le_bytes());
572    hasher.update(item.body.as_bytes());
573    format!("{:x}", hasher.finalize())
574}
575
576pub(crate) fn hash_text(text: &str) -> String {
577    let mut hasher = Sha256::new();
578    hasher.update(text.as_bytes());
579    format!("{:x}", hasher.finalize())
580}
context_harness/app_store.rs

context_harness/
app_store.rs
