1use std::collections::HashSet;
4use std::sync::{Arc, OnceLock};
5
6use async_trait::async_trait;
7use uuid::Uuid;
8
9use khive_score::DeterministicScore;
10use khive_storage::error::StorageError;
11use khive_storage::types::{
12 BatchWriteSummary, IndexRebuildScope, OrphanSweepConfig, OrphanSweepResult, SqlStatement,
13 SqlValue, VectorIndexKind, VectorRecord, VectorSearchHit, VectorSearchRequest,
14 VectorStoreCapabilities, VectorStoreInfo,
15};
16use khive_storage::StorageCapability;
17use khive_storage::StorageResult;
18use khive_storage::VectorStore;
19use khive_types::SubstrateKind;
20
21use crate::error::SqliteError;
22use crate::pool::ConnectionPool;
23use crate::sql_bridge::bind_params;
24
25pub fn delete_vector_statement(table: &str, subject_id: Uuid, namespace: &str) -> SqlStatement {
30 SqlStatement {
31 sql: format!("DELETE FROM {table} WHERE subject_id = ?1 AND namespace = ?2"),
32 params: vec![
33 SqlValue::Text(subject_id.to_string()),
34 SqlValue::Text(namespace.to_string()),
35 ],
36 label: Some(format!("vec-delete-{table}")),
37 }
38}
39
40#[cfg(test)]
51mod failpoint {
52 use std::sync::atomic::{AtomicBool, Ordering};
53 use std::sync::Arc;
54
55 use std::cell::RefCell;
56
57 thread_local! {
58 pub(super) static CURRENT: RefCell<Option<Arc<AtomicBool>>> = const { RefCell::new(None) };
62 }
63
64 #[cfg(feature = "vectors")]
76 pub(super) fn arm() {
77 let flag = Arc::new(AtomicBool::new(true));
78 CURRENT.with(|c| *c.borrow_mut() = Some(flag));
79 }
80
81 #[cfg(feature = "vectors")]
84 pub(super) fn disarm() {
85 CURRENT.with(|c| *c.borrow_mut() = None);
86 }
87
88 pub(super) fn take(flag: &Arc<AtomicBool>) -> bool {
91 flag.compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
92 .is_ok()
93 }
94
95 #[cfg(feature = "vectors")]
99 pub(super) struct FailpointGuard;
100
101 #[cfg(feature = "vectors")]
102 impl FailpointGuard {
103 pub(super) fn new() -> Self {
104 arm();
105 Self
106 }
107 }
108
109 #[cfg(feature = "vectors")]
110 impl Drop for FailpointGuard {
111 fn drop(&mut self) {
112 disarm();
113 }
114 }
115}
116
117fn f32_slice_as_bytes(data: &[f32]) -> &[u8] {
124 unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, std::mem::size_of_val(data)) }
129}
130
131#[cfg(test)]
138fn current_failpoint() -> Option<std::sync::Arc<std::sync::atomic::AtomicBool>> {
139 failpoint::CURRENT.with(|c| c.borrow().clone())
140}
141
142#[cfg(not(test))]
143fn current_failpoint() -> Option<std::sync::Arc<std::sync::atomic::AtomicBool>> {
144 None
145}
146
147fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
148 StorageError::driver(StorageCapability::Vectors, op, e)
149}
150
151fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
152 StorageError::driver(StorageCapability::Vectors, op, e)
153}
154
155fn non_finite_index(data: &[f32]) -> Option<usize> {
156 data.iter().position(|v| !v.is_finite())
157}
158
159fn non_finite_vector_error(op: &'static str, idx: usize, value: f32) -> StorageError {
160 StorageError::InvalidInput {
161 capability: StorageCapability::Vectors,
162 operation: op.into(),
163 message: format!(
164 "non-finite value at index {idx}: {value} \
165 (NaN/Inf values corrupt distance computations)"
166 ),
167 }
168}
169
170fn validate_model_key(model_key: &str) -> Result<(), SqliteError> {
172 if model_key.is_empty()
173 || !model_key
174 .chars()
175 .all(|c| c.is_ascii_alphanumeric() || c == '_')
176 {
177 return Err(SqliteError::InvalidData(format!(
178 "invalid model_key '{}': must be non-empty and contain only ASCII alphanumeric/underscore characters",
179 model_key
180 )));
181 }
182 Ok(())
183}
184
185pub struct SqliteVecStore {
191 pool: Arc<ConnectionPool>,
192 is_file_backed: bool,
193 model_key: String,
194 embedding_model: String,
195 dimensions: usize,
196 table_name: String,
197 namespace: String,
198 writer_task: Option<crate::writer_task::WriterTaskHandle>,
199}
200
201impl SqliteVecStore {
202 pub fn new(
206 pool: Arc<ConnectionPool>,
207 is_file_backed: bool,
208 model_key: String,
209 embedding_model: String,
210 dimensions: usize,
211 namespace: String,
212 ) -> Result<Self, SqliteError> {
213 validate_model_key(&model_key)?;
214 let table_name = format!("vec_{}", model_key);
215 let writer_task = pool.writer_task_handle().ok().flatten();
219 Ok(Self {
220 pool,
221 is_file_backed,
222 model_key,
223 embedding_model,
224 dimensions,
225 table_name,
226 namespace,
227 writer_task,
228 })
229 }
230
231 fn open_standalone_reader(&self) -> Result<rusqlite::Connection, StorageError> {
232 let config = self.pool.config();
233 let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
234 operation: "vec_reader".into(),
235 message: "in-memory databases do not support standalone connections".into(),
236 })?;
237
238 let conn = rusqlite::Connection::open_with_flags(
239 path,
240 rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY
241 | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
242 | rusqlite::OpenFlags::SQLITE_OPEN_URI,
243 )
244 .map_err(|e| map_err(e, "open_vec_reader"))?;
245
246 conn.busy_timeout(config.busy_timeout)
247 .map_err(|e| map_err(e, "open_vec_reader"))?;
248 conn.pragma_update(None, "foreign_keys", "ON")
249 .map_err(|e| map_err(e, "open_vec_reader"))?;
250 conn.pragma_update(None, "synchronous", "NORMAL")
251 .map_err(|e| map_err(e, "open_vec_reader"))?;
252
253 Ok(conn)
254 }
255
256 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
260 where
261 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
262 R: Send + 'static,
263 {
264 if let Some(writer_task) = &self.writer_task {
265 return writer_task
266 .send(move |conn| f(conn).map_err(|e| map_err(e, op)))
267 .await;
268 }
269
270 self.with_writer_unmanaged(op, f).await
271 }
272
273 async fn with_writer_unmanaged<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
278 where
279 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
280 R: Send + 'static,
281 {
282 let pool = Arc::clone(&self.pool);
283 tokio::task::spawn_blocking(move || {
284 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
285 f(guard.conn()).map_err(|e| map_err(e, op))
286 })
287 .await
288 .map_err(|e| StorageError::driver(StorageCapability::Vectors, op, e))?
289 }
290
291 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
292 where
293 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
294 R: Send + 'static,
295 {
296 if self.is_file_backed {
297 let conn = self.open_standalone_reader()?;
298 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
299 .await
300 .map_err(|e| StorageError::driver(StorageCapability::Vectors, op, e))?
301 } else {
302 let pool = Arc::clone(&self.pool);
303 tokio::task::spawn_blocking(move || {
304 let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
305 f(guard.conn()).map_err(|e| map_err(e, op))
306 })
307 .await
308 .map_err(|e| StorageError::driver(StorageCapability::Vectors, op, e))?
309 }
310 }
311}
312
313struct VectorRowRef<'a> {
317 subject_id: Uuid,
318 namespace: &'a str,
319 kind: &'a str,
320 field: &'a str,
321 embedding_model: &'a str,
322 embedding: &'a [f32],
323}
324
325fn replace_vector_row_dml(
329 conn: &rusqlite::Connection,
330 table: &str,
331 dims: usize,
332 row: VectorRowRef<'_>,
333 failpoint_flag: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
334) -> Result<(), rusqlite::Error> {
335 if row.embedding.len() != dims {
336 return Err(rusqlite::Error::InvalidParameterCount(
337 row.embedding.len(),
338 dims,
339 ));
340 }
341
342 let del_sql = format!("DELETE FROM {table} WHERE subject_id = ?1 AND namespace = ?2");
343 conn.execute(
344 &del_sql,
345 rusqlite::params![row.subject_id.to_string(), row.namespace],
346 )?;
347
348 #[cfg(test)]
352 if let Some(ref fp) = failpoint_flag {
353 if failpoint::take(fp) {
354 return Err(rusqlite::Error::InvalidParameterName(
355 "__test_failpoint_after_delete__".into(),
356 ));
357 }
358 }
359 #[cfg(not(test))]
360 let _ = failpoint_flag;
361
362 let ins_sql = format!(
363 "INSERT INTO {table} (subject_id, namespace, kind, field, embedding_model, embedding) \
364 VALUES (?1, ?2, ?3, ?4, ?5, ?6)"
365 );
366 let blob = f32_slice_as_bytes(row.embedding);
367 conn.execute(
368 &ins_sql,
369 rusqlite::params![
370 row.subject_id.to_string(),
371 row.namespace,
372 row.kind,
373 row.field,
374 row.embedding_model,
375 blob
376 ],
377 )?;
378
379 conn.execute(
382 "INSERT INTO ann_write_log (namespace, embedding_model, kind, field, subject_id, op) \
383 VALUES (?1, ?2, ?3, ?4, ?5, 'upsert')",
384 rusqlite::params![
385 row.namespace,
386 row.embedding_model,
387 row.kind,
388 row.field,
389 row.subject_id.to_string()
390 ],
391 )?;
392
393 Ok(())
394}
395
396fn log_vector_deletes(
401 conn: &rusqlite::Connection,
402 table: &str,
403 where_clause: &str,
404 params: &[&dyn rusqlite::ToSql],
405) -> Result<(), rusqlite::Error> {
406 let sql = format!(
407 "INSERT INTO ann_write_log (namespace, embedding_model, kind, field, subject_id, op) \
408 SELECT namespace, embedding_model, kind, field, subject_id, 'delete' \
409 FROM {table} WHERE {where_clause}"
410 );
411 conn.execute(&sql, params)?;
412 Ok(())
413}
414
415pub fn delete_subject_from_vector_tables(
423 conn: &rusqlite::Connection,
424 tables: &[String],
425 subject_id: Uuid,
426 namespace: &str,
427) -> Result<(), rusqlite::Error> {
428 for table in tables {
429 log_vector_deletes(
430 conn,
431 table,
432 "subject_id = ?1 AND namespace = ?2",
433 &[&subject_id.to_string(), &namespace],
434 )?;
435 let sql = format!("DELETE FROM {table} WHERE subject_id = ?1 AND namespace = ?2");
436 conn.execute(&sql, rusqlite::params![subject_id.to_string(), namespace])?;
437 }
438 Ok(())
439}
440
441#[allow(clippy::too_many_arguments)]
450fn batch_insert_vectors_dml(
451 conn: &rusqlite::Connection,
452 table: &str,
453 dims: usize,
454 store_embedding_model: &str,
455 records: &[VectorRecord],
456 attempted: u64,
457 failpoint_flag: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
458) -> Result<BatchWriteSummary, rusqlite::Error> {
459 let mut affected = 0u64;
460 let mut failed = 0u64;
461 let mut first_error = String::new();
462
463 for record in records {
464 if record.vectors.len() != 1 {
465 if first_error.is_empty() {
466 first_error = format!("expected 1 vector per record, got {}", record.vectors.len());
467 }
468 failed += 1;
469 continue;
470 }
471 let embedding = &record.vectors[0];
472 if embedding.len() != dims {
473 if first_error.is_empty() {
474 first_error = format!(
475 "wrong vector dimension: expected {dims}, got {}",
476 embedding.len()
477 );
478 }
479 failed += 1;
480 continue;
481 }
482 if non_finite_index(embedding).is_some() {
483 if first_error.is_empty() {
484 first_error = "embedding contains non-finite values (NaN or Inf)".to_string();
485 }
486 failed += 1;
487 continue;
488 }
489 let kind_str = record.kind.to_string();
490
491 conn.execute_batch("SAVEPOINT vec_batch_record")?;
495 let result = replace_vector_row_dml(
496 conn,
497 table,
498 dims,
499 VectorRowRef {
500 subject_id: record.subject_id,
501 namespace: &record.namespace,
502 kind: &kind_str,
503 field: &record.field,
504 embedding_model: store_embedding_model,
505 embedding,
506 },
507 failpoint_flag.clone(),
508 );
509 match result {
510 Ok(()) => {
511 conn.execute_batch("RELEASE SAVEPOINT vec_batch_record")?;
512 affected += 1;
513 }
514 Err(e) => {
515 let _ = conn.execute_batch("ROLLBACK TO SAVEPOINT vec_batch_record");
516 let _ = conn.execute_batch("RELEASE SAVEPOINT vec_batch_record");
517 if first_error.is_empty() {
518 first_error = e.to_string();
519 }
520 failed += 1;
521 }
522 }
523 }
524
525 Ok(BatchWriteSummary {
526 attempted,
527 affected,
528 failed,
529 first_error,
530 })
531}
532
533#[allow(clippy::too_many_arguments)]
542fn vec_upsert_atomic_dml(
543 conn: &rusqlite::Connection,
544 table: &str,
545 dims: usize,
546 subject_id: Uuid,
547 kind_str: &str,
548 namespace: &str,
549 field: &str,
550 embedding_model: &str,
551 embedding: &[f32],
552 savepoint_name: &'static str,
553 failpoint_flag: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
554) -> Result<(), rusqlite::Error> {
555 conn.execute_batch(&format!("SAVEPOINT {savepoint_name}"))?;
556 let result = replace_vector_row_dml(
557 conn,
558 table,
559 dims,
560 VectorRowRef {
561 subject_id,
562 namespace,
563 kind: kind_str,
564 field,
565 embedding_model,
566 embedding,
567 },
568 failpoint_flag,
569 );
570
571 match result {
572 Ok(()) => {
573 conn.execute_batch(&format!("RELEASE SAVEPOINT {savepoint_name}"))?;
574 Ok(())
575 }
576 Err(e) => {
577 let _ = conn.execute_batch(&format!("ROLLBACK TO SAVEPOINT {savepoint_name}"));
578 let _ = conn.execute_batch(&format!("RELEASE SAVEPOINT {savepoint_name}"));
579 Err(e)
580 }
581 }
582}
583
584fn orphan_sweep_dml(
593 conn: &rusqlite::Connection,
594 table: &str,
595 ns_json: Option<&str>,
596 kind_json: Option<&str>,
597 allow_json: Option<&str>,
598 max_delete: i64,
599 dry_run: bool,
600) -> Result<OrphanSweepResult, rusqlite::Error> {
601 let filter_pred = "(?1 IS NULL OR namespace IN (SELECT value FROM json_each(?1))) \
607 AND (?2 IS NULL OR kind IN (SELECT value FROM json_each(?2))) \
608 AND (?3 IS NULL OR subject_id IN (SELECT value FROM json_each(?3)))";
609
610 let live_subq = "SELECT id FROM entities WHERE deleted_at IS NULL \
619 UNION ALL \
620 SELECT id FROM notes WHERE deleted_at IS NULL";
621
622 let orphan_pred = format!(
623 "subject_id NOT IN ({live}) AND {f}",
624 live = live_subq,
625 f = filter_pred,
626 );
627
628 let scan_sql = format!(
630 "SELECT COUNT(*) FROM {t} WHERE {f}",
631 t = table,
632 f = filter_pred
633 );
634 let scanned: i64 = conn.query_row(
635 &scan_sql,
636 rusqlite::params![ns_json, kind_json, allow_json],
637 |row| row.get(0),
638 )?;
639
640 let count_sql = format!(
642 "SELECT COUNT(*) FROM {t} WHERE {p}",
643 t = table,
644 p = orphan_pred,
645 );
646 let would_delete: i64 = conn.query_row(
647 &count_sql,
648 rusqlite::params![ns_json, kind_json, allow_json],
649 |row| row.get(0),
650 )?;
651
652 let max_delete_hit = would_delete > max_delete;
653
654 let deleted: i64 = if dry_run {
665 0
666 } else {
667 let select_sql = format!(
668 "SELECT subject_id FROM {t} WHERE {p} LIMIT ?4",
669 t = table,
670 p = orphan_pred,
671 );
672 let mut stmt = conn.prepare(&select_sql)?;
673 let victim_ids: Vec<String> = stmt
674 .query_map(
675 rusqlite::params![ns_json, kind_json, allow_json, max_delete],
676 |row| row.get::<_, String>(0),
677 )?
678 .collect::<Result<_, _>>()?;
679 drop(stmt);
680
681 let mut total: i64 = 0;
682 for chunk in victim_ids.chunks(400) {
683 let placeholders: String = (1..=chunk.len())
684 .map(|i| format!("?{i}"))
685 .collect::<Vec<_>>()
686 .join(", ");
687 let in_clause = format!("subject_id IN ({placeholders})");
688 let params: Vec<&dyn rusqlite::ToSql> =
689 chunk.iter().map(|s| s as &dyn rusqlite::ToSql).collect();
690 log_vector_deletes(conn, table, &in_clause, ¶ms)?;
691 let del_sql = format!("DELETE FROM {t} WHERE {in_clause}", t = table);
692 let mut del_stmt = conn.prepare(&del_sql)?;
693 for (i, id_str) in chunk.iter().enumerate() {
694 del_stmt.raw_bind_parameter(i + 1, id_str.as_str())?;
695 }
696 total += del_stmt.raw_execute()? as i64;
697 }
698 total
699 };
700
701 Ok(OrphanSweepResult {
702 scanned: scanned as u64,
703 would_delete: would_delete as u64,
704 deleted: deleted as u64,
705 max_delete_hit,
706 })
707}
708
709#[async_trait]
710impl VectorStore for SqliteVecStore {
711 async fn insert(
712 &self,
713 subject_id: Uuid,
714 kind: SubstrateKind,
715 namespace: &str,
716 field: &str,
717 vectors: Vec<Vec<f32>>,
718 ) -> Result<(), StorageError> {
719 if vectors.len() != 1 {
720 return Err(StorageError::Unsupported {
721 capability: StorageCapability::Vectors,
722 operation: "vec_insert".into(),
723 message: "sqlite-vec supports exactly one vector per record".into(),
724 });
725 }
726 let embedding = vectors.into_iter().next().expect("len checked");
727
728 let table = self.table_name.clone();
729 let dims = self.dimensions;
730 let namespace = namespace.to_string();
731 let field = field.to_string();
732 let kind_str = kind.to_string();
733 let embedding_model = self.embedding_model.clone();
734
735 if embedding.len() == dims {
736 if let Some(idx) = non_finite_index(&embedding) {
737 return Err(non_finite_vector_error("vec_insert", idx, embedding[idx]));
738 }
739 }
740
741 let failpoint_flag = current_failpoint();
744
745 if let Some(writer_task) = &self.writer_task {
752 let table2 = table.clone();
753 let namespace2 = namespace.clone();
754 let field2 = field.clone();
755 let kind_str2 = kind_str.clone();
756 let embedding_model2 = embedding_model.clone();
757 let embedding2 = embedding.clone();
758 return writer_task
759 .send(move |conn| {
760 vec_upsert_atomic_dml(
761 conn,
762 &table2,
763 dims,
764 subject_id,
765 &kind_str2,
766 &namespace2,
767 &field2,
768 &embedding_model2,
769 &embedding2,
770 "vec_insert_atomic",
771 failpoint_flag,
772 )
773 .map_err(|e| map_err(e, "vec_insert"))
774 })
775 .await;
776 }
777
778 let origin = self.pool.origin();
783 self.with_writer("vec_insert", move |conn| {
784 let _tx_handle = khive_storage::tx_registry::register_scoped(
790 Some("vec_insert_tx".to_string()),
791 origin,
792 );
793 let tx = conn.unchecked_transaction()?;
794
795 replace_vector_row_dml(
796 &tx,
797 &table,
798 dims,
799 VectorRowRef {
800 subject_id,
801 namespace: &namespace,
802 kind: &kind_str,
803 field: &field,
804 embedding_model: &embedding_model,
805 embedding: &embedding,
806 },
807 failpoint_flag,
808 )?;
809
810 tx.commit()
811 })
812 .await
813 }
814
815 async fn insert_batch(
816 &self,
817 records: Vec<VectorRecord>,
818 ) -> Result<BatchWriteSummary, StorageError> {
819 let table = self.table_name.clone();
820 let dims = self.dimensions;
821 let attempted = records.len() as u64;
822 let store_embedding_model = self.embedding_model.clone();
823
824 let failpoint_flag = current_failpoint();
829
830 if let Some(writer_task) = &self.writer_task {
836 let table2 = table.clone();
837 let store_embedding_model2 = store_embedding_model.clone();
838 return writer_task
839 .send(move |conn| {
840 batch_insert_vectors_dml(
841 conn,
842 &table2,
843 dims,
844 &store_embedding_model2,
845 &records,
846 attempted,
847 failpoint_flag,
848 )
849 .map_err(|e| map_err(e, "vec_insert_batch"))
850 })
851 .await;
852 }
853
854 let origin = self.pool.origin();
857 self.with_writer("vec_insert_batch", move |conn| {
858 conn.execute_batch("BEGIN IMMEDIATE")?;
859 let _tx_handle = khive_storage::tx_registry::register_scoped(
860 Some("vector_insert_batch".to_string()),
861 origin,
862 );
863
864 let summary = batch_insert_vectors_dml(
865 conn,
866 &table,
867 dims,
868 &store_embedding_model,
869 &records,
870 attempted,
871 failpoint_flag,
872 )?;
873
874 conn.execute_batch("COMMIT")?;
875
876 Ok(summary)
877 })
878 .await
879 }
880
881 async fn update(
882 &self,
883 subject_id: Uuid,
884 kind: SubstrateKind,
885 namespace: &str,
886 field: &str,
887 vectors: Vec<Vec<f32>>,
888 ) -> Result<(), StorageError> {
889 if vectors.len() != 1 {
890 return Err(StorageError::Unsupported {
891 capability: StorageCapability::Vectors,
892 operation: "vec_update".into(),
893 message: "sqlite-vec supports exactly one vector per record".into(),
894 });
895 }
896 let embedding = vectors.into_iter().next().expect("len checked");
897
898 let table = self.table_name.clone();
899 let dims = self.dimensions;
900 let namespace = namespace.to_string();
901 let field = field.to_string();
902 let kind_str = kind.to_string();
903 let embedding_model = self.embedding_model.clone();
904
905 if embedding.len() == dims {
906 if let Some(idx) = non_finite_index(&embedding) {
907 return Err(non_finite_vector_error("vec_update", idx, embedding[idx]));
908 }
909 }
910
911 let failpoint_flag = current_failpoint();
914
915 if let Some(writer_task) = &self.writer_task {
922 let table2 = table.clone();
923 let namespace2 = namespace.clone();
924 let field2 = field.clone();
925 let kind_str2 = kind_str.clone();
926 let embedding_model2 = embedding_model.clone();
927 let embedding2 = embedding.clone();
928 return writer_task
929 .send(move |conn| {
930 vec_upsert_atomic_dml(
931 conn,
932 &table2,
933 dims,
934 subject_id,
935 &kind_str2,
936 &namespace2,
937 &field2,
938 &embedding_model2,
939 &embedding2,
940 "vec_update_atomic",
941 failpoint_flag,
942 )
943 .map_err(|e| map_err(e, "vec_update"))
944 })
945 .await;
946 }
947
948 let origin = self.pool.origin();
952 self.with_writer("vec_update", move |conn| {
953 let _tx_handle = khive_storage::tx_registry::register_scoped(
956 Some("vec_update_tx".to_string()),
957 origin,
958 );
959 let tx = conn.unchecked_transaction()?;
960
961 replace_vector_row_dml(
962 &tx,
963 &table,
964 dims,
965 VectorRowRef {
966 subject_id,
967 namespace: &namespace,
968 kind: &kind_str,
969 field: &field,
970 embedding_model: &embedding_model,
971 embedding: &embedding,
972 },
973 failpoint_flag,
974 )?;
975
976 tx.commit()
977 })
978 .await
979 }
980
981 async fn delete(&self, subject_id: Uuid) -> Result<bool, StorageError> {
982 let statement = delete_vector_statement(&self.table_name, subject_id, &self.namespace);
983 let table = self.table_name.clone();
984 let namespace = self.namespace.clone();
985
986 self.with_writer("vec_delete", move |conn| {
987 conn.execute_batch("SAVEPOINT vec_delete_log")?;
988 let result = (|| {
989 log_vector_deletes(
990 conn,
991 &table,
992 "subject_id = ?1 AND namespace = ?2",
993 &[&subject_id.to_string(), &namespace],
994 )?;
995 let mut stmt = conn.prepare(&statement.sql)?;
996 bind_params(&mut stmt, &statement.params)?;
997 Ok(stmt.raw_execute()? > 0)
998 })();
999 match result {
1000 Ok(v) => {
1001 conn.execute_batch("RELEASE SAVEPOINT vec_delete_log")?;
1002 Ok(v)
1003 }
1004 Err(e) => {
1005 let _ = conn.execute_batch("ROLLBACK TO SAVEPOINT vec_delete_log");
1006 let _ = conn.execute_batch("RELEASE SAVEPOINT vec_delete_log");
1007 Err(e)
1008 }
1009 }
1010 })
1011 .await
1012 }
1013
1014 async fn count(&self) -> Result<u64, StorageError> {
1015 let table = self.table_name.clone();
1016 let namespace = self.namespace.clone();
1017
1018 self.with_reader("vec_count", move |conn| {
1019 let sql = format!("SELECT COUNT(*) FROM {} WHERE namespace = ?1", table);
1020 let count: i64 =
1021 conn.query_row(&sql, rusqlite::params![&namespace], |row| row.get(0))?;
1022 Ok(count as u64)
1023 })
1024 .await
1025 }
1026
1027 async fn search(
1028 &self,
1029 request: VectorSearchRequest,
1030 ) -> Result<Vec<VectorSearchHit>, StorageError> {
1031 if request.filter.as_ref().is_some_and(|f| !f.is_empty()) {
1032 return Err(StorageError::Unsupported {
1033 capability: StorageCapability::Vectors,
1034 operation: "vec_search".into(),
1035 message: "use search_with_filter for filtered queries".into(),
1036 });
1037 }
1038 if request.query_vectors.len() != 1 {
1039 return Err(StorageError::Unsupported {
1040 capability: StorageCapability::Vectors,
1041 operation: "vec_search".into(),
1042 message: "sqlite-vec supports exactly one query vector per search".into(),
1043 });
1044 }
1045 let query_embedding = request.query_vectors[0].clone();
1046
1047 let table = self.table_name.clone();
1048 let dims = self.dimensions;
1049 let namespace = request
1051 .namespace
1052 .clone()
1053 .unwrap_or_else(|| self.namespace.clone());
1054 let kind_filter = request.kind.map(|k| k.to_string());
1055 let effective_model = request
1057 .embedding_model
1058 .clone()
1059 .unwrap_or_else(|| self.embedding_model.clone());
1060
1061 if query_embedding.len() == dims {
1062 if let Some(idx) = non_finite_index(&query_embedding) {
1063 return Err(non_finite_vector_error(
1064 "vec_search",
1065 idx,
1066 query_embedding[idx],
1067 ));
1068 }
1069 }
1070
1071 self.with_reader("vec_search", move |conn| {
1072 if query_embedding.len() != dims {
1073 return Err(rusqlite::Error::InvalidParameterCount(
1074 query_embedding.len(),
1075 dims,
1076 ));
1077 }
1078
1079 let kind_clause = if kind_filter.is_some() {
1083 "AND kind = ?5"
1084 } else {
1085 ""
1086 };
1087 let sql = format!(
1088 "SELECT subject_id, distance \
1089 FROM {t} \
1090 WHERE embedding MATCH ?1 \
1091 AND namespace = ?3 \
1092 AND embedding_model = ?4 \
1093 {kind_clause} \
1094 ORDER BY distance \
1095 LIMIT ?2",
1096 t = table,
1097 kind_clause = kind_clause
1098 );
1099
1100 let query_blob = f32_slice_as_bytes(&query_embedding);
1101 let mut stmt = conn.prepare(&sql)?;
1102
1103 let raw_rows: Vec<rusqlite::Result<(String, f64)>> =
1106 if let Some(ref kind_str) = kind_filter {
1107 stmt.query_map(
1108 rusqlite::params![
1109 query_blob,
1110 request.top_k,
1111 &namespace,
1112 &effective_model,
1113 kind_str
1114 ],
1115 |row| {
1116 let id_str: String = row.get(0)?;
1117 let distance: f64 = row.get(1)?;
1118 Ok((id_str, distance))
1119 },
1120 )?
1121 .collect()
1122 } else {
1123 stmt.query_map(
1124 rusqlite::params![query_blob, request.top_k, &namespace, &effective_model],
1125 |row| {
1126 let id_str: String = row.get(0)?;
1127 let distance: f64 = row.get(1)?;
1128 Ok((id_str, distance))
1129 },
1130 )?
1131 .collect()
1132 };
1133
1134 let mut hits = Vec::new();
1135 for (rank_idx, row) in raw_rows.into_iter().enumerate() {
1136 let (id_str, distance) = row?;
1137 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1138 rusqlite::Error::FromSqlConversionFailure(
1139 0,
1140 rusqlite::types::Type::Text,
1141 Box::new(e),
1142 )
1143 })?;
1144
1145 let similarity = 1.0 - (distance / 2.0);
1148
1149 hits.push(VectorSearchHit {
1150 subject_id,
1151 score: DeterministicScore::from_f64(similarity),
1152 rank: (rank_idx + 1) as u32,
1153 });
1154 }
1155
1156 Ok(hits)
1157 })
1158 .await
1159 }
1160
1161 async fn info(&self) -> Result<VectorStoreInfo, StorageError> {
1162 let count = self.count().await?;
1163
1164 Ok(VectorStoreInfo {
1165 model_name: self.model_key.clone(),
1166 dimensions: self.dimensions,
1167 index_kind: VectorIndexKind::SqliteVec,
1168 entry_count: count,
1169 needs_rebuild: false,
1170 last_rebuild_at: None,
1171 })
1172 }
1173
1174 async fn rebuild(&self, _scope: IndexRebuildScope) -> Result<VectorStoreInfo, StorageError> {
1175 self.info().await
1177 }
1178
1179 async fn delete_subjects(&self, ids: &[Uuid]) -> Result<u64, StorageError> {
1180 if ids.is_empty() {
1181 return Ok(0);
1182 }
1183 let table = self.table_name.clone();
1184 let id_strings: Vec<String> = ids.iter().map(|id| id.to_string()).collect();
1185 let mut total_deleted: u64 = 0;
1186
1187 for chunk in id_strings.chunks(400) {
1189 let placeholders: String = (1..=chunk.len())
1190 .map(|i| format!("?{i}"))
1191 .collect::<Vec<_>>()
1192 .join(", ");
1193 let sql = format!("DELETE FROM {table} WHERE subject_id IN ({placeholders})");
1194 let in_clause = format!("subject_id IN ({placeholders})");
1195 let chunk_owned = chunk.to_vec();
1196 let table_cl = table.clone();
1197 let table_sp = table.clone();
1198 let deleted = self
1199 .with_writer("vec_delete_subjects", move |conn| {
1200 conn.execute_batch("SAVEPOINT vec_delete_subjects_log")?;
1201 let result = (|| {
1202 let params: Vec<&dyn rusqlite::ToSql> = chunk_owned
1203 .iter()
1204 .map(|s| s as &dyn rusqlite::ToSql)
1205 .collect();
1206 log_vector_deletes(conn, &table_sp, &in_clause, ¶ms)?;
1207 let mut stmt = conn.prepare(&sql)?;
1208 for (i, id_str) in chunk_owned.iter().enumerate() {
1209 stmt.raw_bind_parameter(i + 1, id_str.as_str())?;
1210 }
1211 stmt.raw_execute().map(|n| n as u64)
1212 })();
1213 match result {
1214 Ok(n) => {
1215 conn.execute_batch("RELEASE SAVEPOINT vec_delete_subjects_log")?;
1216 Ok(n)
1217 }
1218 Err(e) => {
1219 let _ =
1220 conn.execute_batch("ROLLBACK TO SAVEPOINT vec_delete_subjects_log");
1221 let _ = conn.execute_batch("RELEASE SAVEPOINT vec_delete_subjects_log");
1222 Err(e)
1223 }
1224 }
1225 })
1226 .await
1227 .map_err(|e| {
1228 tracing::warn!(error = %e, table = %table_cl, "delete_subjects chunk failed");
1229 e
1230 })?;
1231 total_deleted += deleted;
1232 }
1233 Ok(total_deleted)
1234 }
1235
1236 async fn batch_exists(
1237 &self,
1238 ids: &[Uuid],
1239 namespace: &str,
1240 ) -> Result<HashSet<Uuid>, StorageError> {
1241 if ids.is_empty() {
1242 return Ok(HashSet::new());
1243 }
1244
1245 let table = self.table_name.clone();
1246 let namespace = namespace.to_string();
1247 let model = self.embedding_model.clone();
1248 let id_strings: Vec<String> = ids.iter().map(|id| id.to_string()).collect();
1249
1250 self.with_reader("vec_batch_exists", move |conn| {
1251 let mut found = HashSet::new();
1252
1253 for chunk in id_strings.chunks(400) {
1254 let placeholders: String = (0..chunk.len())
1256 .map(|i| format!("?{}", i + 3))
1257 .collect::<Vec<_>>()
1258 .join(", ");
1259
1260 let sql = format!(
1261 "SELECT subject_id FROM {} WHERE namespace = ?1 \
1262 AND embedding_model = ?2 AND subject_id IN ({})",
1263 table, placeholders
1264 );
1265
1266 let mut stmt = conn.prepare(&sql)?;
1267 stmt.raw_bind_parameter(1, namespace.as_str())?;
1268 stmt.raw_bind_parameter(2, model.as_str())?;
1269 for (i, id_str) in chunk.iter().enumerate() {
1270 stmt.raw_bind_parameter(i + 3, id_str.as_str())?;
1271 }
1272
1273 let mut rows = stmt.raw_query();
1274 while let Some(row) = rows.next()? {
1275 let id_str: String = row.get(0)?;
1276 if let Ok(uuid) = Uuid::parse_str(&id_str) {
1277 found.insert(uuid);
1278 }
1279 }
1280 }
1281
1282 Ok(found)
1283 })
1284 .await
1285 }
1286
1287 async fn orphan_sweep(&self, config: &OrphanSweepConfig) -> StorageResult<OrphanSweepResult> {
1288 let table = self.table_name.clone();
1289
1290 let ns_json: Option<String> = if config.namespaces.is_empty() {
1294 None
1295 } else {
1296 serde_json::to_string(&config.namespaces).ok()
1297 };
1298
1299 let kind_json: Option<String> = if config.substrate_kinds.is_empty() {
1300 None
1301 } else {
1302 let strs: Vec<String> = config
1303 .substrate_kinds
1304 .iter()
1305 .map(|k| k.to_string())
1306 .collect();
1307 serde_json::to_string(&strs).ok()
1308 };
1309
1310 let allow_json: Option<String> = config.subject_id_allowlist.as_ref().map(|ids| {
1312 let strs: Vec<String> = ids.iter().map(|id| id.to_string()).collect();
1313 serde_json::to_string(&strs).unwrap_or_default()
1314 });
1315
1316 let max_delete = config.max_delete as i64;
1317 let dry_run = config.dry_run;
1318
1319 if let Some(writer_task) = &self.writer_task {
1328 let table2 = table.clone();
1329 let ns_json2 = ns_json.clone();
1330 let kind_json2 = kind_json.clone();
1331 let allow_json2 = allow_json.clone();
1332 return writer_task
1333 .send(move |conn| {
1334 orphan_sweep_dml(
1335 conn,
1336 &table2,
1337 ns_json2.as_deref(),
1338 kind_json2.as_deref(),
1339 allow_json2.as_deref(),
1340 max_delete,
1341 dry_run,
1342 )
1343 .map_err(|e| map_err(e, "orphan_sweep"))
1344 })
1345 .await;
1346 }
1347
1348 let origin = self.pool.origin();
1352 self.with_writer_unmanaged("orphan_sweep", move |conn| {
1353 let _tx_handle = khive_storage::tx_registry::register_scoped(
1370 Some("vec_orphan_sweep".to_string()),
1371 origin,
1372 );
1373 let tx = rusqlite::Transaction::new_unchecked(
1374 conn,
1375 rusqlite::TransactionBehavior::Immediate,
1376 )?;
1377
1378 let result = orphan_sweep_dml(
1379 conn,
1380 &table,
1381 ns_json.as_deref(),
1382 kind_json.as_deref(),
1383 allow_json.as_deref(),
1384 max_delete,
1385 dry_run,
1386 )?;
1387
1388 tx.commit()?;
1389
1390 Ok(result)
1391 })
1392 .await
1393 }
1394
1395 fn capabilities(&self) -> &'static VectorStoreCapabilities {
1396 static SQLITE_VEC_CAPABILITIES: OnceLock<VectorStoreCapabilities> = OnceLock::new();
1397 SQLITE_VEC_CAPABILITIES.get_or_init(|| VectorStoreCapabilities {
1398 supports_filter: false,
1399 supports_batch_search: false,
1400 supports_quantization: false,
1401 supports_update: false,
1402 supports_orphan_sweep: true,
1403 supports_multi_field: false,
1408 max_dimensions: Some(8192),
1413 index_kinds: vec![VectorIndexKind::SqliteVec],
1414 })
1415 }
1416}
1417
1418impl SqliteVecStore {
1419 pub async fn score_candidates(
1425 &self,
1426 query_embedding: &[f32],
1427 candidate_ids: &[Uuid],
1428 ) -> Result<Vec<VectorSearchHit>, StorageError> {
1429 if candidate_ids.is_empty() || query_embedding.is_empty() {
1430 return Ok(Vec::new());
1431 }
1432
1433 let dims = self.dimensions;
1434 if query_embedding.len() != dims {
1435 return Err(StorageError::InvalidInput {
1436 capability: StorageCapability::Vectors,
1437 operation: "score_candidates".into(),
1438 message: format!(
1439 "query has {} dims, expected {}",
1440 query_embedding.len(),
1441 dims
1442 ),
1443 });
1444 }
1445
1446 if let Some(idx) = non_finite_index(query_embedding) {
1447 return Err(non_finite_vector_error(
1448 "score_candidates",
1449 idx,
1450 query_embedding[idx],
1451 ));
1452 }
1453
1454 let table = self.table_name.clone();
1455 let namespace = self.namespace.clone();
1456 let embedding_model = self.embedding_model.clone();
1457 let query_vec = query_embedding.to_vec();
1458 let ids: Vec<String> = candidate_ids.iter().map(|id| id.to_string()).collect();
1459
1460 self.with_reader("score_candidates", move |conn| {
1461 let mut all_hits: Vec<VectorSearchHit> = Vec::new();
1462 let query_blob = f32_slice_as_bytes(&query_vec);
1463
1464 for chunk in ids.chunks(399) {
1465 let placeholders: String = chunk
1466 .iter()
1467 .enumerate()
1468 .map(|(i, _)| format!("?{}", i + 4))
1469 .collect::<Vec<_>>()
1470 .join(", ");
1471
1472 let sql = format!(
1473 "SELECT e.subject_id, vec_distance_cosine(e.embedding, ?1) as distance \
1474 FROM {} e \
1475 WHERE e.namespace = ?2 AND e.embedding_model = ?3 \
1476 AND e.subject_id IN ({})",
1477 table, placeholders
1478 );
1479
1480 let mut stmt = conn.prepare(&sql)?;
1481 stmt.raw_bind_parameter(1, query_blob)?;
1482 stmt.raw_bind_parameter(2, namespace.as_str())?;
1483 stmt.raw_bind_parameter(3, embedding_model.as_str())?;
1484 for (i, id_str) in chunk.iter().enumerate() {
1485 stmt.raw_bind_parameter(i + 4, id_str.as_str())?;
1486 }
1487
1488 let mut rows = stmt.raw_query();
1489 while let Some(row) = rows.next()? {
1490 let id_str: String = row.get(0)?;
1491 let distance: f64 = row.get(1)?;
1492
1493 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1494 rusqlite::Error::FromSqlConversionFailure(
1495 0,
1496 rusqlite::types::Type::Text,
1497 Box::new(e),
1498 )
1499 })?;
1500
1501 let similarity = 1.0 - (distance / 2.0);
1502 all_hits.push(VectorSearchHit {
1503 subject_id,
1504 score: DeterministicScore::from_f64(similarity),
1505 rank: 0,
1506 });
1507 }
1508 }
1509
1510 all_hits.sort_by_key(|hit| std::cmp::Reverse(hit.score));
1511 for (i, hit) in all_hits.iter_mut().enumerate() {
1512 hit.rank = (i + 1) as u32;
1513 }
1514
1515 Ok(all_hits)
1516 })
1517 .await
1518 }
1519}
1520
1521#[cfg(all(test, feature = "vectors"))]
1522mod batch_exists_tests {
1523 use std::collections::HashSet;
1524 use std::sync::Arc;
1525
1526 use khive_types::SubstrateKind;
1527 use uuid::Uuid;
1528
1529 use super::*;
1530
1531 fn make_vec_pool() -> Arc<crate::pool::ConnectionPool> {
1532 use crate::pool::{ConnectionPool, PoolConfig};
1533 crate::extension::ensure_extensions_loaded();
1534 let config = PoolConfig {
1535 path: None,
1536 ..PoolConfig::default()
1537 };
1538 Arc::new(ConnectionPool::new(config).expect("in-memory pool"))
1539 }
1540
1541 fn create_vec_table(pool: &Arc<crate::pool::ConnectionPool>, model_key: &str, dims: usize) {
1542 let writer = pool.try_writer().expect("pool writer");
1543 let ddl = format!(
1544 "CREATE VIRTUAL TABLE IF NOT EXISTS vec_{} USING vec0(\
1545 subject_id TEXT PRIMARY KEY, \
1546 namespace TEXT NOT NULL, \
1547 kind TEXT NOT NULL, \
1548 field TEXT NOT NULL, \
1549 embedding_model TEXT NOT NULL, \
1550 embedding float[{}] distance_metric=cosine)",
1551 model_key, dims
1552 );
1553 writer.conn().execute_batch(&ddl).expect("create vec table");
1554 writer
1555 .conn()
1556 .execute_batch(crate::migrations::ANN_WRITE_LOG_DDL)
1557 .expect("create ann_write_log");
1558 }
1559
1560 #[tokio::test]
1563 async fn batch_exists_returns_correct_set_for_underscored_model_key() {
1564 let pool = make_vec_pool();
1565 let model_key = "all_minilm_l6_v2";
1566 let dims = 4;
1567 let ns = "ns:test";
1568
1569 create_vec_table(&pool, model_key, dims);
1570
1571 let store = SqliteVecStore::new(
1572 pool,
1573 false,
1574 model_key.to_string(),
1575 model_key.to_string(),
1576 dims,
1577 ns.to_string(),
1578 )
1579 .expect("SqliteVecStore::new");
1580
1581 let id1 = Uuid::new_v4();
1582 let id2 = Uuid::new_v4();
1583 let id_absent = Uuid::new_v4();
1584
1585 store
1586 .insert(
1587 id1,
1588 SubstrateKind::Entity,
1589 ns,
1590 "body",
1591 vec![vec![0.1, 0.2, 0.3, 0.4]],
1592 )
1593 .await
1594 .expect("insert id1");
1595 store
1596 .insert(
1597 id2,
1598 SubstrateKind::Entity,
1599 ns,
1600 "body",
1601 vec![vec![0.5, 0.6, 0.7, 0.8]],
1602 )
1603 .await
1604 .expect("insert id2");
1605
1606 let exists = store
1607 .batch_exists(&[id1, id2, id_absent], ns)
1608 .await
1609 .expect("batch_exists");
1610
1611 assert!(exists.contains(&id1), "id1 must be found");
1612 assert!(exists.contains(&id2), "id2 must be found");
1613 assert!(
1614 !exists.contains(&id_absent),
1615 "absent id must not be returned"
1616 );
1617 assert_eq!(exists.len(), 2);
1618 }
1619
1620 #[tokio::test]
1622 async fn batch_exists_empty_ids_returns_empty_set() {
1623 let pool = make_vec_pool();
1624 let model_key = "empty_test_model";
1625 create_vec_table(&pool, model_key, 4);
1626
1627 let store = SqliteVecStore::new(
1628 pool,
1629 false,
1630 model_key.to_string(),
1631 model_key.to_string(),
1632 4,
1633 "ns:test".to_string(),
1634 )
1635 .expect("SqliteVecStore::new");
1636
1637 let exists: HashSet<Uuid> = store
1638 .batch_exists(&[], "ns:test")
1639 .await
1640 .expect("batch_exists");
1641 assert!(exists.is_empty());
1642 }
1643
1644 #[tokio::test]
1649 async fn vector_search_namespace_predicate_prevents_recall_starvation() {
1650 let pool = make_vec_pool();
1651 let model_key = "knn_namespace_scope";
1652 let dims = 4;
1653 create_vec_table(&pool, model_key, dims);
1654
1655 let store = SqliteVecStore::new(
1656 pool,
1657 false,
1658 model_key.to_string(),
1659 model_key.to_string(),
1660 dims,
1661 "ns:b".to_string(),
1662 )
1663 .expect("SqliteVecStore::new");
1664
1665 let distractor_a = Uuid::new_v4();
1666 let victim_b = Uuid::new_v4();
1667
1668 store
1670 .insert(
1671 distractor_a,
1672 SubstrateKind::Entity,
1673 "ns:a",
1674 "body",
1675 vec![vec![1.0, 0.0, 0.0, 0.0]],
1676 )
1677 .await
1678 .expect("insert nearer cross-namespace vector");
1679
1680 store
1682 .insert(
1683 victim_b,
1684 SubstrateKind::Entity,
1685 "ns:b",
1686 "body",
1687 vec![vec![0.8, 0.2, 0.0, 0.0]],
1688 )
1689 .await
1690 .expect("insert in-namespace vector");
1691
1692 let hits = store
1694 .search(VectorSearchRequest {
1695 query_vectors: vec![vec![1.0, 0.0, 0.0, 0.0]],
1696 top_k: 1,
1697 namespace: Some("ns:b".to_string()),
1698 kind: Some(SubstrateKind::Entity),
1699 embedding_model: None,
1700 filter: None,
1701 backend_hints: None,
1702 })
1703 .await
1704 .expect("search");
1705
1706 assert_eq!(
1707 hits.len(),
1708 1,
1709 "namespace B must not be starved by namespace A"
1710 );
1711 assert_eq!(
1712 hits[0].subject_id, victim_b,
1713 "top-1 in ns:b must be victim_b, not cross-namespace distractor_a"
1714 );
1715 }
1716
1717 #[test]
1720 fn hyphenated_model_key_is_rejected_at_construction() {
1721 use crate::pool::{ConnectionPool, PoolConfig};
1722 let pool = Arc::new(
1723 ConnectionPool::new(PoolConfig {
1724 path: None,
1725 ..PoolConfig::default()
1726 })
1727 .expect("pool"),
1728 );
1729
1730 let result = SqliteVecStore::new(
1731 pool,
1732 false,
1733 "all-minilm-l6-v2".to_string(),
1734 "all-minilm-l6-v2".to_string(),
1735 4,
1736 "ns:test".to_string(),
1737 );
1738
1739 assert!(
1740 result.is_err(),
1741 "hyphenated model_key 'all-minilm-l6-v2' must be rejected; \
1742 the store's table_name would differ from what a hand-rolled sanitizer produces"
1743 );
1744 }
1745}
1746
1747#[cfg(test)]
1753mod first_error_tests {
1754 use super::*;
1755 use khive_storage::types::VectorRecord;
1756 use khive_storage::VectorStore;
1757 use khive_types::SubstrateKind;
1758 use uuid::Uuid;
1759
1760 fn make_pool() -> Arc<crate::pool::ConnectionPool> {
1761 use crate::pool::{ConnectionPool, PoolConfig};
1762 let config = PoolConfig {
1763 path: None,
1764 ..PoolConfig::default()
1765 };
1766 Arc::new(ConnectionPool::new(config).expect("in-memory pool"))
1767 }
1768
1769 #[tokio::test]
1780 async fn insert_batch_first_error_populated_on_dimension_mismatch() {
1781 let dims = 4usize;
1782 let store = SqliteVecStore::new(
1783 make_pool(),
1784 false,
1785 "first_err_vec".into(),
1786 "first_err_vec".into(),
1787 dims,
1788 "ns:test".into(),
1789 )
1790 .expect("SqliteVecStore::new");
1791
1792 let summary = store
1795 .insert_batch(vec![
1796 VectorRecord {
1797 subject_id: Uuid::new_v4(),
1798 kind: SubstrateKind::Entity,
1799 namespace: "ns:test".to_string(),
1800 field: "body".to_string(),
1801 embedding_model: None,
1802 vectors: vec![vec![0.0f32; dims + 1]],
1803 updated_at: chrono::Utc::now(),
1804 },
1805 VectorRecord {
1806 subject_id: Uuid::new_v4(),
1807 kind: SubstrateKind::Entity,
1808 namespace: "ns:test".to_string(),
1809 field: "body".to_string(),
1810 embedding_model: None,
1811 vectors: vec![vec![0.0f32; dims + 2]],
1812 updated_at: chrono::Utc::now(),
1813 },
1814 ])
1815 .await
1816 .expect("insert_batch must return Ok (best-effort semantics)");
1817
1818 assert_eq!(summary.attempted, 2);
1819 assert_eq!(
1820 summary.failed, 2,
1821 "both wrong-dims records must be counted as failed"
1822 );
1823 assert_eq!(summary.affected, 0);
1824 assert!(
1825 !summary.first_error.is_empty(),
1826 "first_error must be populated when failed > 0; \
1827 got empty string; the validation error is silently swallowed"
1828 );
1829 }
1830}
1831
1832#[cfg(test)]
1833mod capabilities_tests {
1834 use super::*;
1835
1836 fn make_pool() -> Arc<crate::pool::ConnectionPool> {
1837 use crate::pool::{ConnectionPool, PoolConfig};
1838 let config = PoolConfig {
1839 path: None,
1840 ..PoolConfig::default()
1841 };
1842 Arc::new(ConnectionPool::new(config).expect("in-memory pool"))
1843 }
1844
1845 #[test]
1846 fn sqlite_vec_store_capabilities_are_correct() {
1847 let store = SqliteVecStore::new(
1848 make_pool(),
1849 false,
1850 "test_model".into(),
1851 "test_model".into(),
1852 4,
1853 "ns:test".into(),
1854 )
1855 .expect("SqliteVecStore::new");
1856
1857 let caps = store.capabilities();
1858
1859 assert!(
1860 !caps.supports_filter,
1861 "sqlite-vec does not support filter pushdown"
1862 );
1863 assert!(
1864 !caps.supports_batch_search,
1865 "sqlite-vec does not support native batch search"
1866 );
1867 assert!(
1868 !caps.supports_quantization,
1869 "sqlite-vec does not support quantization"
1870 );
1871 assert!(
1872 !caps.supports_update,
1873 "sqlite-vec does not support in-place update"
1874 );
1875 assert!(
1876 caps.supports_orphan_sweep,
1877 "SqliteVecStore must advertise supports_orphan_sweep = true"
1878 );
1879 assert_eq!(caps.max_dimensions, Some(8192));
1881 assert_eq!(
1882 caps.index_kinds,
1883 vec![VectorIndexKind::SqliteVec],
1884 "index_kinds should be [SqliteVec]"
1885 );
1886 }
1887
1888 #[test]
1892 fn max_dimensions_reflects_sqlite_vec_hard_limit_not_k_max() {
1893 let store = SqliteVecStore::new(
1894 make_pool(),
1895 false,
1896 "test_dim_limit".into(),
1897 "test_dim_limit".into(),
1898 4,
1899 "ns:test".into(),
1900 )
1901 .expect("SqliteVecStore::new");
1902
1903 let caps = store.capabilities();
1904
1905 let max = caps
1909 .max_dimensions
1910 .expect("SqliteVecStore must declare a finite dimension limit");
1911 assert!(
1912 max >= 8192,
1913 "max_dimensions ({max}) must be at least 8192 — the sqlite-vec hard limit"
1914 );
1915 }
1916
1917 #[test]
1920 fn capabilities_is_idempotent() {
1921 let store = SqliteVecStore::new(
1922 make_pool(),
1923 false,
1924 "test_idempotent".into(),
1925 "test_idempotent".into(),
1926 4,
1927 "ns:test".into(),
1928 )
1929 .expect("SqliteVecStore::new");
1930
1931 let caps1 = store.capabilities();
1932 let caps2 = store.capabilities();
1933 assert_eq!(
1934 caps1 as *const _, caps2 as *const _,
1935 "capabilities() must return the same static reference each call"
1936 );
1937 }
1938}
1939
1940#[cfg(all(test, feature = "vectors"))]
1941mod atomic_replace_tests {
1942 use std::sync::Arc;
1943
1944 use khive_storage::types::VectorRecord;
1945 use khive_storage::VectorStore;
1946 use khive_types::SubstrateKind;
1947 use uuid::Uuid;
1948
1949 use super::*;
1950
1951 fn make_vec_pool() -> Arc<crate::pool::ConnectionPool> {
1952 use crate::pool::{ConnectionPool, PoolConfig};
1953 crate::extension::ensure_extensions_loaded();
1954 let config = PoolConfig {
1955 path: None,
1956 ..PoolConfig::default()
1957 };
1958 Arc::new(ConnectionPool::new(config).expect("in-memory pool"))
1959 }
1960
1961 fn create_vec_table(pool: &Arc<crate::pool::ConnectionPool>, model_key: &str, dims: usize) {
1962 let writer = pool.try_writer().expect("pool writer");
1963 let ddl = format!(
1964 "CREATE VIRTUAL TABLE IF NOT EXISTS vec_{} USING vec0(\
1965 subject_id TEXT PRIMARY KEY, \
1966 namespace TEXT NOT NULL, \
1967 kind TEXT NOT NULL, \
1968 field TEXT NOT NULL, \
1969 embedding_model TEXT NOT NULL, \
1970 embedding float[{}] distance_metric=cosine)",
1971 model_key, dims
1972 );
1973 writer.conn().execute_batch(&ddl).expect("create vec table");
1974 writer
1975 .conn()
1976 .execute_batch(crate::migrations::ANN_WRITE_LOG_DDL)
1977 .expect("create ann_write_log");
1978 }
1979
1980 #[tokio::test]
1989 async fn insert_batch_failed_record_preserves_prior_vector() {
1990 let pool = make_vec_pool();
1991 let model_key = "atomic_batch_test";
1992 let dims = 4;
1993 let ns = "ns:atomic";
1994
1995 create_vec_table(&pool, model_key, dims);
1996
1997 let store = SqliteVecStore::new(
1998 Arc::clone(&pool),
1999 false,
2000 model_key.to_string(),
2001 model_key.to_string(),
2002 dims,
2003 ns.to_string(),
2004 )
2005 .expect("SqliteVecStore::new");
2006
2007 let id_existing = Uuid::new_v4();
2008 let id_new = Uuid::new_v4();
2009 let original_vec = vec![0.1f32, 0.2, 0.3, 0.4];
2010
2011 store
2012 .insert(
2013 id_existing,
2014 SubstrateKind::Entity,
2015 ns,
2016 "body",
2017 vec![original_vec.clone()],
2018 )
2019 .await
2020 .expect("initial insert");
2021
2022 let summary = store
2023 .insert_batch(vec![
2024 VectorRecord {
2025 subject_id: id_existing,
2026 kind: SubstrateKind::Entity,
2027 namespace: ns.to_string(),
2028 field: "body".to_string(),
2029 embedding_model: None,
2030 vectors: vec![vec![9.9f32; dims + 1]],
2031 updated_at: chrono::Utc::now(),
2032 },
2033 VectorRecord {
2034 subject_id: id_new,
2035 kind: SubstrateKind::Entity,
2036 namespace: ns.to_string(),
2037 field: "body".to_string(),
2038 embedding_model: None,
2039 vectors: vec![vec![0.5f32, 0.6, 0.7, 0.8]],
2040 updated_at: chrono::Utc::now(),
2041 },
2042 ])
2043 .await
2044 .expect("insert_batch");
2045
2046 assert_eq!(summary.attempted, 2);
2047 assert_eq!(summary.affected, 1, "only id_new should succeed");
2048 assert_eq!(summary.failed, 1, "id_existing with wrong dims must fail");
2049
2050 let existing_still_present = store
2051 .batch_exists(&[id_existing], ns)
2052 .await
2053 .expect("batch_exists");
2054 assert!(
2055 existing_still_present.contains(&id_existing),
2056 "prior vector for id_existing must survive a failed batch replace"
2057 );
2058
2059 let new_present = store
2060 .batch_exists(&[id_new], ns)
2061 .await
2062 .expect("batch_exists for id_new");
2063 assert!(
2064 new_present.contains(&id_new),
2065 "id_new with valid dims must be inserted"
2066 );
2067 }
2068
2069 #[tokio::test]
2072 async fn update_failed_preserves_prior_vector() {
2073 let pool = make_vec_pool();
2074 let model_key = "atomic_update_test";
2075 let dims = 4;
2076 let ns = "ns:atomic_upd";
2077
2078 create_vec_table(&pool, model_key, dims);
2079
2080 let store = SqliteVecStore::new(
2081 Arc::clone(&pool),
2082 false,
2083 model_key.to_string(),
2084 model_key.to_string(),
2085 dims,
2086 ns.to_string(),
2087 )
2088 .expect("SqliteVecStore::new");
2089
2090 let id = Uuid::new_v4();
2091
2092 store
2093 .insert(
2094 id,
2095 SubstrateKind::Entity,
2096 ns,
2097 "body",
2098 vec![vec![0.1f32, 0.2, 0.3, 0.4]],
2099 )
2100 .await
2101 .expect("initial insert");
2102
2103 let result = store
2104 .update(
2105 id,
2106 SubstrateKind::Entity,
2107 ns,
2108 "body",
2109 vec![vec![9.9f32; dims + 1]],
2110 )
2111 .await;
2112
2113 assert!(result.is_err(), "update with wrong dims must fail");
2114
2115 let still_present = store
2116 .batch_exists(&[id], ns)
2117 .await
2118 .expect("batch_exists after failed update");
2119 assert!(
2120 still_present.contains(&id),
2121 "prior vector must survive a failed update"
2122 );
2123 }
2124
2125 #[tokio::test]
2129 async fn insert_batch_savepoint_rollback_on_pk_conflict_preserves_stale() {
2130 let pool = make_vec_pool();
2131 let model_key = "atomic_pk_batch";
2132 let dims = 4;
2133 let ns_a = "ns:pk_a";
2134 let ns_b = "ns:pk_b";
2135
2136 create_vec_table(&pool, model_key, dims);
2137
2138 let store = SqliteVecStore::new(
2139 Arc::clone(&pool),
2140 false,
2141 model_key.to_string(),
2142 model_key.to_string(),
2143 dims,
2144 ns_a.to_string(),
2145 )
2146 .expect("SqliteVecStore::new");
2147
2148 let id_x = Uuid::new_v4();
2149 let stale_vec = vec![0.1f32, 0.2, 0.3, 0.4];
2150
2151 store
2153 .insert(
2154 id_x,
2155 SubstrateKind::Entity,
2156 ns_a,
2157 "body",
2158 vec![stale_vec.clone()],
2159 )
2160 .await
2161 .expect("stale insert");
2162
2163 let summary = store
2168 .insert_batch(vec![VectorRecord {
2169 subject_id: id_x,
2170 kind: SubstrateKind::Entity,
2171 namespace: ns_b.to_string(),
2172 field: "body".to_string(),
2173 embedding_model: None,
2174 vectors: vec![vec![0.5f32, 0.6, 0.7, 0.8]],
2175 updated_at: chrono::Utc::now(),
2176 }])
2177 .await
2178 .expect("insert_batch must complete (outer tx must commit)");
2179
2180 assert_eq!(summary.attempted, 1);
2181 assert_eq!(summary.affected, 0, "PK conflict must count as failed");
2182 assert_eq!(
2183 summary.failed, 1,
2184 "failed counter must increment after ROLLBACK TO SAVEPOINT"
2185 );
2186
2187 let post = store
2189 .batch_exists(&[id_x], ns_a)
2190 .await
2191 .expect("batch_exists ns:a");
2192 assert!(
2193 post.contains(&id_x),
2194 "stale row in ns:a must survive after SAVEPOINT + INSERT failure"
2195 );
2196
2197 let hits = store
2200 .search(VectorSearchRequest {
2201 query_vectors: vec![stale_vec.clone()],
2202 top_k: 1,
2203 namespace: Some(ns_a.to_string()),
2204 kind: Some(SubstrateKind::Entity),
2205 embedding_model: None,
2206 filter: None,
2207 backend_hints: None,
2208 })
2209 .await
2210 .expect("search ns:a after batch");
2211
2212 assert_eq!(hits.len(), 1, "stale vector must be searchable");
2213 assert_eq!(hits[0].subject_id, id_x);
2214 let sim = hits[0].score.to_f64();
2215 assert!(
2216 sim > 0.999,
2217 "cosine similarity of stale_vec to itself must be ~1.0 (got {sim:.6}); \
2218 a lower value means the SAVEPOINT/ROLLBACK left partial writes visible"
2219 );
2220 }
2221
2222 #[tokio::test]
2226 async fn insert_batch_rollback_does_not_corrupt_subsequent_record() {
2227 let pool = make_vec_pool();
2228 let model_key = "atomic_sib_batch";
2229 let dims = 4;
2230 let ns_a = "ns:sib_a";
2231 let ns_b = "ns:sib_b";
2232
2233 create_vec_table(&pool, model_key, dims);
2234
2235 let store = SqliteVecStore::new(
2236 Arc::clone(&pool),
2237 false,
2238 model_key.to_string(),
2239 model_key.to_string(),
2240 dims,
2241 ns_a.to_string(),
2242 )
2243 .expect("SqliteVecStore::new");
2244
2245 let id_x = Uuid::new_v4();
2246 let stale_vec = vec![0.1f32, 0.2, 0.3, 0.4];
2247 let new_vec = vec![0.9f32, 0.1, 0.1, 0.1];
2248
2249 store
2251 .insert(
2252 id_x,
2253 SubstrateKind::Entity,
2254 ns_a,
2255 "body",
2256 vec![stale_vec.clone()],
2257 )
2258 .await
2259 .expect("stale insert");
2260
2261 let summary = store
2263 .insert_batch(vec![
2264 VectorRecord {
2265 subject_id: id_x,
2266 kind: SubstrateKind::Entity,
2267 namespace: ns_b.to_string(),
2268 field: "body".to_string(),
2269 embedding_model: None,
2270 vectors: vec![vec![0.5f32, 0.6, 0.7, 0.8]],
2271 updated_at: chrono::Utc::now(),
2272 },
2273 VectorRecord {
2274 subject_id: id_x,
2275 kind: SubstrateKind::Entity,
2276 namespace: ns_a.to_string(),
2277 field: "body".to_string(),
2278 embedding_model: None,
2279 vectors: vec![new_vec.clone()],
2280 updated_at: chrono::Utc::now(),
2281 },
2282 ])
2283 .await
2284 .expect("insert_batch");
2285
2286 assert_eq!(summary.attempted, 2);
2287 assert_eq!(summary.affected, 1, "Record B must succeed");
2290 assert_eq!(summary.failed, 1, "Record A must fail (PK conflict)");
2291
2292 let hits = store
2294 .search(VectorSearchRequest {
2295 query_vectors: vec![new_vec.clone()],
2296 top_k: 1,
2297 namespace: Some(ns_a.to_string()),
2298 kind: Some(SubstrateKind::Entity),
2299 embedding_model: None,
2300 filter: None,
2301 backend_hints: None,
2302 })
2303 .await
2304 .expect("search after batch");
2305
2306 assert_eq!(hits.len(), 1);
2307 assert_eq!(hits[0].subject_id, id_x);
2308 let sim = hits[0].score.to_f64();
2309 assert!(
2310 sim > 0.999,
2311 "new_vec similarity to itself must be ~1.0 (got {sim:.6}); \
2312 Record A's ROLLBACK must not corrupt Record B's write"
2313 );
2314 }
2315
2316 #[tokio::test]
2320 async fn update_pk_conflict_rolls_back_transaction_preserves_stale() {
2321 let pool = make_vec_pool();
2322 let model_key = "atomic_upd_pk";
2323 let dims = 4;
2324 let ns_a = "ns:upk_a";
2325 let ns_b = "ns:upk_b";
2326
2327 create_vec_table(&pool, model_key, dims);
2328
2329 let store = SqliteVecStore::new(
2330 Arc::clone(&pool),
2331 false,
2332 model_key.to_string(),
2333 model_key.to_string(),
2334 dims,
2335 ns_a.to_string(),
2336 )
2337 .expect("store");
2338
2339 let id_x = Uuid::new_v4();
2340 let stale_vec = vec![0.1f32, 0.2, 0.3, 0.4];
2341
2342 store
2344 .insert(
2345 id_x,
2346 SubstrateKind::Entity,
2347 ns_a,
2348 "body",
2349 vec![stale_vec.clone()],
2350 )
2351 .await
2352 .expect("stale insert");
2353
2354 let result = store
2357 .update(
2358 id_x,
2359 SubstrateKind::Entity,
2360 ns_b,
2361 "body",
2362 vec![vec![0.5f32, 0.6, 0.7, 0.8]],
2363 )
2364 .await;
2365
2366 assert!(
2367 result.is_err(),
2368 "update must fail when INSERT hits the vec0 PK constraint"
2369 );
2370
2371 let post = store
2373 .batch_exists(&[id_x], ns_a)
2374 .await
2375 .expect("batch_exists after failed update");
2376 assert!(
2377 post.contains(&id_x),
2378 "stale row in ns:a must survive after update transaction rollback"
2379 );
2380
2381 let hits = store
2383 .search(VectorSearchRequest {
2384 query_vectors: vec![stale_vec.clone()],
2385 top_k: 1,
2386 namespace: Some(ns_a.to_string()),
2387 kind: Some(SubstrateKind::Entity),
2388 embedding_model: None,
2389 filter: None,
2390 backend_hints: None,
2391 })
2392 .await
2393 .expect("search after failed update");
2394
2395 assert_eq!(hits.len(), 1, "stale vector must be searchable");
2396 assert_eq!(hits[0].subject_id, id_x);
2397 let sim = hits[0].score.to_f64();
2398 assert!(
2399 sim > 0.999,
2400 "cosine similarity of stale_vec to itself must be ~1.0 (got {sim:.6}); \
2401 transaction rollback must leave embedding bytes unchanged"
2402 );
2403 }
2404
2405 #[tokio::test]
2412 async fn insert_batch_rollback_restores_deleted_stale_after_post_delete_insert_failure() {
2413 let pool = make_vec_pool();
2414 let model_key = "sentinel_batch_rb";
2415 let dims = 4;
2416 let ns = "ns:sentinel_batch";
2417
2418 create_vec_table(&pool, model_key, dims);
2419
2420 let store = SqliteVecStore::new(
2421 Arc::clone(&pool),
2422 false,
2423 model_key.to_string(),
2424 model_key.to_string(),
2425 dims,
2426 ns.to_string(),
2427 )
2428 .expect("SqliteVecStore::new");
2429
2430 let id_x = Uuid::new_v4();
2431 let vec1 = vec![0.1f32, 0.2, 0.3, 0.4];
2432 let vec2 = vec![0.9f32, 0.0, 0.0, 0.0];
2433
2434 store
2436 .insert(id_x, SubstrateKind::Entity, ns, "body", vec![vec1.clone()])
2437 .await
2438 .expect("stale insert");
2439
2440 let _guard = failpoint::FailpointGuard::new();
2444
2445 let summary = store
2447 .insert_batch(vec![VectorRecord {
2448 subject_id: id_x,
2449 kind: SubstrateKind::Entity,
2450 namespace: ns.to_string(),
2451 field: "body".to_string(),
2452 embedding_model: None,
2453 vectors: vec![vec2.clone()],
2454 updated_at: chrono::Utc::now(),
2455 }])
2456 .await
2457 .expect("insert_batch must complete (outer tx must commit regardless)");
2458
2459 drop(_guard); assert_eq!(summary.attempted, 1);
2462 assert_eq!(
2463 summary.affected, 0,
2464 "failpoint must prevent INSERT from succeeding"
2465 );
2466 assert_eq!(
2467 summary.failed, 1,
2468 "failed counter must increment after injected failure"
2469 );
2470
2471 let present = store
2473 .batch_exists(&[id_x], ns)
2474 .await
2475 .expect("batch_exists after failpoint");
2476 assert!(
2477 present.contains(&id_x),
2478 "ROLLBACK TO SAVEPOINT must restore the stale row after DELETE + injected failure"
2479 );
2480
2481 let hits = store
2483 .search(VectorSearchRequest {
2484 query_vectors: vec![vec1.clone()],
2485 top_k: 1,
2486 namespace: Some(ns.to_string()),
2487 kind: Some(SubstrateKind::Entity),
2488 embedding_model: None,
2489 filter: None,
2490 backend_hints: None,
2491 })
2492 .await
2493 .expect("search after failpoint");
2494
2495 assert_eq!(
2496 hits.len(),
2497 1,
2498 "stale vector must be searchable after rollback"
2499 );
2500 assert_eq!(hits[0].subject_id, id_x);
2501 let sim = hits[0].score.to_f64();
2502 assert!(
2503 sim > 0.999,
2504 "similarity to vec1 must be ~1.0 (got {sim:.6}); \
2505 a lower value means the stale embedding was not restored — ROLLBACK TO SAVEPOINT failed"
2506 );
2507
2508 let hits2 = store
2510 .search(VectorSearchRequest {
2511 query_vectors: vec![vec2.clone()],
2512 top_k: 1,
2513 namespace: Some(ns.to_string()),
2514 kind: Some(SubstrateKind::Entity),
2515 embedding_model: None,
2516 filter: None,
2517 backend_hints: None,
2518 })
2519 .await
2520 .expect("search vec2 after failpoint");
2521 let sim2 = hits2.first().map(|h| h.score.to_f64()).unwrap_or(0.0);
2522 assert!(
2523 sim2 < 0.99,
2524 "similarity to vec2 must be < 0.99 (got {sim2:.6}); \
2525 vec2 must not be the stored embedding after a rolled-back INSERT"
2526 );
2527 }
2528
2529 #[tokio::test]
2533 async fn update_rollback_restores_deleted_stale_after_post_delete_insert_failure() {
2534 let pool = make_vec_pool();
2535 let model_key = "sentinel_upd_rb";
2536 let dims = 4;
2537 let ns = "ns:sentinel_upd";
2538
2539 create_vec_table(&pool, model_key, dims);
2540
2541 let store = SqliteVecStore::new(
2542 Arc::clone(&pool),
2543 false,
2544 model_key.to_string(),
2545 model_key.to_string(),
2546 dims,
2547 ns.to_string(),
2548 )
2549 .expect("SqliteVecStore::new");
2550
2551 let id_x = Uuid::new_v4();
2552 let vec1 = vec![0.1f32, 0.2, 0.3, 0.4];
2553 let vec2 = vec![0.9f32, 0.0, 0.0, 0.0];
2554
2555 store
2557 .insert(id_x, SubstrateKind::Entity, ns, "body", vec![vec1.clone()])
2558 .await
2559 .expect("stale insert");
2560
2561 let _guard = failpoint::FailpointGuard::new();
2563
2564 let result = store
2566 .update(id_x, SubstrateKind::Entity, ns, "body", vec![vec2.clone()])
2567 .await;
2568
2569 drop(_guard);
2570
2571 assert!(
2572 result.is_err(),
2573 "update must propagate the injected error back to the caller"
2574 );
2575
2576 let present = store
2578 .batch_exists(&[id_x], ns)
2579 .await
2580 .expect("batch_exists after failpoint");
2581 assert!(
2582 present.contains(&id_x),
2583 "transaction rollback must restore the stale row after DELETE + injected failure"
2584 );
2585
2586 let hits = store
2588 .search(VectorSearchRequest {
2589 query_vectors: vec![vec1.clone()],
2590 top_k: 1,
2591 namespace: Some(ns.to_string()),
2592 kind: Some(SubstrateKind::Entity),
2593 embedding_model: None,
2594 filter: None,
2595 backend_hints: None,
2596 })
2597 .await
2598 .expect("search after failpoint");
2599
2600 assert_eq!(
2601 hits.len(),
2602 1,
2603 "stale vector must be searchable after rollback"
2604 );
2605 assert_eq!(hits[0].subject_id, id_x);
2606 let sim = hits[0].score.to_f64();
2607 assert!(
2608 sim > 0.999,
2609 "similarity to vec1 must be ~1.0 (got {sim:.6}); \
2610 a lower value means the stale embedding was not restored — transaction rollback failed"
2611 );
2612 }
2613
2614 #[tokio::test]
2619 async fn insert_rollback_restores_deleted_stale_after_post_delete_insert_failure() {
2620 let pool = make_vec_pool();
2621 let model_key = "sentinel_ins_rb";
2622 let dims = 4;
2623 let ns = "ns:sentinel_ins";
2624
2625 create_vec_table(&pool, model_key, dims);
2626
2627 let store = SqliteVecStore::new(
2628 Arc::clone(&pool),
2629 false,
2630 model_key.to_string(),
2631 model_key.to_string(),
2632 dims,
2633 ns.to_string(),
2634 )
2635 .expect("SqliteVecStore::new");
2636
2637 let id_x = Uuid::new_v4();
2638 let vec1 = vec![0.1f32, 0.2, 0.3, 0.4];
2639 let vec2 = vec![0.9f32, 0.0, 0.0, 0.0];
2640
2641 store
2645 .insert(id_x, SubstrateKind::Entity, ns, "body", vec![vec1.clone()])
2646 .await
2647 .expect("stale insert");
2648
2649 let _guard = failpoint::FailpointGuard::new();
2651
2652 let result = store
2654 .insert(id_x, SubstrateKind::Entity, ns, "body", vec![vec2.clone()])
2655 .await;
2656
2657 drop(_guard);
2658
2659 assert!(
2660 result.is_err(),
2661 "insert must propagate the injected error back to the caller"
2662 );
2663
2664 let present = store
2666 .batch_exists(&[id_x], ns)
2667 .await
2668 .expect("batch_exists after failpoint");
2669 assert!(
2670 present.contains(&id_x),
2671 "transaction rollback must restore the stale row after DELETE + injected failure"
2672 );
2673
2674 let hits = store
2676 .search(VectorSearchRequest {
2677 query_vectors: vec![vec1.clone()],
2678 top_k: 1,
2679 namespace: Some(ns.to_string()),
2680 kind: Some(SubstrateKind::Entity),
2681 embedding_model: None,
2682 filter: None,
2683 backend_hints: None,
2684 })
2685 .await
2686 .expect("search after failpoint");
2687
2688 assert_eq!(
2689 hits.len(),
2690 1,
2691 "stale vector must be searchable after rollback"
2692 );
2693 assert_eq!(hits[0].subject_id, id_x);
2694 let sim = hits[0].score.to_f64();
2695 assert!(
2696 sim > 0.999,
2697 "similarity to vec1 must be ~1.0 (got {sim:.6}); \
2698 a lower value means the stale embedding was not restored — transaction rollback failed"
2699 );
2700 }
2701}
2702
2703#[cfg(all(test, feature = "vectors"))]
2710mod orphan_sweep_tests {
2711 use std::sync::Arc;
2712
2713 use khive_storage::types::{OrphanSweepConfig, OrphanSweepResult};
2714 use khive_storage::VectorStore;
2715 use khive_types::SubstrateKind;
2716 use uuid::Uuid;
2717
2718 use super::*;
2719
2720 fn make_pool() -> Arc<crate::pool::ConnectionPool> {
2723 use crate::pool::{ConnectionPool, PoolConfig};
2724 crate::extension::ensure_extensions_loaded();
2725 Arc::new(
2726 ConnectionPool::new(PoolConfig {
2727 path: None,
2728 ..PoolConfig::default()
2729 })
2730 .expect("in-memory pool"),
2731 )
2732 }
2733
2734 fn create_substrate_tables(pool: &Arc<crate::pool::ConnectionPool>) {
2736 pool.try_writer()
2737 .expect("writer")
2738 .conn()
2739 .execute_batch(
2740 "CREATE TABLE IF NOT EXISTS entities \
2741 (id TEXT PRIMARY KEY, deleted_at INTEGER); \
2742 CREATE TABLE IF NOT EXISTS notes \
2743 (id TEXT PRIMARY KEY, deleted_at INTEGER);",
2744 )
2745 .expect("create substrate tables");
2746 }
2747
2748 fn create_vec_table(pool: &Arc<crate::pool::ConnectionPool>, model_key: &str, dims: usize) {
2749 let ddl = format!(
2750 "CREATE VIRTUAL TABLE IF NOT EXISTS vec_{} USING vec0(\
2751 subject_id TEXT PRIMARY KEY, \
2752 namespace TEXT NOT NULL, \
2753 kind TEXT NOT NULL, \
2754 field TEXT NOT NULL, \
2755 embedding_model TEXT NOT NULL, \
2756 embedding float[{}] distance_metric=cosine)",
2757 model_key, dims
2758 );
2759 let writer = pool.try_writer().expect("writer");
2760 writer.conn().execute_batch(&ddl).expect("create vec table");
2761 writer
2762 .conn()
2763 .execute_batch(crate::migrations::ANN_WRITE_LOG_DDL)
2764 .expect("create ann_write_log");
2765 }
2766
2767 fn make_store(
2768 pool: Arc<crate::pool::ConnectionPool>,
2769 model_key: &str,
2770 dims: usize,
2771 ns: &str,
2772 ) -> SqliteVecStore {
2773 SqliteVecStore::new(
2774 pool,
2775 false,
2776 model_key.to_string(),
2777 model_key.to_string(),
2778 dims,
2779 ns.to_string(),
2780 )
2781 .expect("SqliteVecStore::new")
2782 }
2783
2784 fn insert_entity(pool: &Arc<crate::pool::ConnectionPool>, id: Uuid, deleted_at: Option<i64>) {
2786 let id_str = id.to_string();
2787 pool.try_writer()
2788 .expect("writer")
2789 .conn()
2790 .execute(
2791 "INSERT INTO entities (id, deleted_at) VALUES (?1, ?2)",
2792 rusqlite::params![id_str, deleted_at],
2793 )
2794 .expect("insert entity");
2795 }
2796
2797 fn vec4(a: f32, b: f32, c: f32, d: f32) -> Vec<f32> {
2798 vec![a, b, c, d]
2799 }
2800
2801 fn sweep_all(max_delete: u32, dry_run: bool) -> OrphanSweepConfig {
2802 OrphanSweepConfig {
2803 subject_id_allowlist: None,
2804 namespaces: vec![],
2805 substrate_kinds: vec![],
2806 max_delete,
2807 dry_run,
2808 }
2809 }
2810
2811 #[tokio::test]
2814 async fn orphan_sweep_keeps_live_subject() {
2815 let pool = make_pool();
2816 create_substrate_tables(&pool);
2817 create_vec_table(&pool, "sw_live", 4);
2818 let store = make_store(Arc::clone(&pool), "sw_live", 4, "ns:sw");
2819 let ns = "ns:sw";
2820
2821 let id = Uuid::new_v4();
2822 insert_entity(&pool, id, None); store
2825 .insert(
2826 id,
2827 SubstrateKind::Entity,
2828 ns,
2829 "body",
2830 vec![vec4(0.1, 0.2, 0.3, 0.4)],
2831 )
2832 .await
2833 .expect("insert vec");
2834
2835 let r: OrphanSweepResult = store
2836 .orphan_sweep(&sweep_all(100, false))
2837 .await
2838 .expect("sweep");
2839
2840 assert_eq!(r.scanned, 1, "one vec row exists");
2841 assert_eq!(r.would_delete, 0, "live subject is not an orphan");
2842 assert_eq!(r.deleted, 0);
2843 assert!(!r.max_delete_hit);
2844
2845 let present = store.batch_exists(&[id], ns).await.expect("exists");
2846 assert!(present.contains(&id), "live subject's vec must survive");
2847 }
2848
2849 #[tokio::test]
2852 async fn orphan_sweep_sweeps_soft_deleted_subject() {
2853 let pool = make_pool();
2854 create_substrate_tables(&pool);
2855 create_vec_table(&pool, "sw_soft", 4);
2856 let store = make_store(Arc::clone(&pool), "sw_soft", 4, "ns:soft");
2857 let ns = "ns:soft";
2858
2859 let id = Uuid::new_v4();
2860 insert_entity(&pool, id, Some(1_000_000)); store
2863 .insert(
2864 id,
2865 SubstrateKind::Entity,
2866 ns,
2867 "body",
2868 vec![vec4(0.5, 0.5, 0.5, 0.5)],
2869 )
2870 .await
2871 .expect("insert vec");
2872
2873 let r = store
2874 .orphan_sweep(&sweep_all(100, false))
2875 .await
2876 .expect("sweep");
2877
2878 assert_eq!(r.scanned, 1);
2879 assert_eq!(r.would_delete, 1, "soft-deleted subject counts as orphan");
2880 assert_eq!(r.deleted, 1);
2881 assert!(!r.max_delete_hit);
2882
2883 let present = store.batch_exists(&[id], ns).await.expect("exists");
2884 assert!(
2885 !present.contains(&id),
2886 "soft-deleted subject's vec must be swept"
2887 );
2888 }
2889
2890 #[tokio::test]
2893 async fn orphan_sweep_sweeps_absent_subject() {
2894 let pool = make_pool();
2895 create_substrate_tables(&pool);
2896 create_vec_table(&pool, "sw_absent", 4);
2897 let store = make_store(Arc::clone(&pool), "sw_absent", 4, "ns:absent");
2898 let ns = "ns:absent";
2899
2900 let id = Uuid::new_v4(); store
2903 .insert(
2904 id,
2905 SubstrateKind::Entity,
2906 ns,
2907 "body",
2908 vec![vec4(0.1, 0.2, 0.3, 0.4)],
2909 )
2910 .await
2911 .expect("insert vec");
2912
2913 let r = store
2914 .orphan_sweep(&sweep_all(100, false))
2915 .await
2916 .expect("sweep");
2917
2918 assert_eq!(r.scanned, 1);
2919 assert_eq!(r.would_delete, 1, "absent subject counts as orphan");
2920 assert_eq!(r.deleted, 1);
2921
2922 let present = store.batch_exists(&[id], ns).await.expect("exists");
2923 assert!(!present.contains(&id), "absent subject's vec must be swept");
2924 }
2925
2926 #[tokio::test]
2929 async fn orphan_sweep_dry_run_does_not_delete() {
2930 let pool = make_pool();
2931 create_substrate_tables(&pool);
2932 create_vec_table(&pool, "sw_dry", 4);
2933 let store = make_store(Arc::clone(&pool), "sw_dry", 4, "ns:dry");
2934 let ns = "ns:dry";
2935
2936 let id = Uuid::new_v4(); store
2938 .insert(
2939 id,
2940 SubstrateKind::Entity,
2941 ns,
2942 "body",
2943 vec![vec4(0.1, 0.2, 0.3, 0.4)],
2944 )
2945 .await
2946 .expect("insert vec");
2947
2948 let r = store
2949 .orphan_sweep(&sweep_all(100, true))
2950 .await
2951 .expect("sweep");
2952
2953 assert_eq!(r.would_delete, 1, "dry-run must still count the orphan");
2954 assert_eq!(r.deleted, 0, "dry-run must not delete anything");
2955
2956 let present = store.batch_exists(&[id], ns).await.expect("exists");
2957 assert!(present.contains(&id), "dry-run must not remove the vec");
2958 }
2959
2960 #[tokio::test]
2963 async fn orphan_sweep_max_delete_caps_deletion() {
2964 let pool = make_pool();
2965 create_substrate_tables(&pool);
2966 create_vec_table(&pool, "sw_cap", 4);
2967 let store = make_store(Arc::clone(&pool), "sw_cap", 4, "ns:cap");
2968 let ns = "ns:cap";
2969
2970 let ids: Vec<Uuid> = (0..5).map(|_| Uuid::new_v4()).collect();
2972 for (i, &id) in ids.iter().enumerate() {
2973 let v = i as f32 / 10.0;
2974 store
2975 .insert(
2976 id,
2977 SubstrateKind::Entity,
2978 ns,
2979 "body",
2980 vec![vec![v, v + 0.1, v + 0.2, v + 0.3]],
2981 )
2982 .await
2983 .expect("insert vec");
2984 }
2985
2986 let r = store
2987 .orphan_sweep(&OrphanSweepConfig {
2988 subject_id_allowlist: None,
2989 namespaces: vec![],
2990 substrate_kinds: vec![],
2991 max_delete: 2,
2992 dry_run: false,
2993 })
2994 .await
2995 .expect("sweep");
2996
2997 assert_eq!(r.scanned, 5);
2998 assert_eq!(r.would_delete, 5);
2999 assert_eq!(r.deleted, 2, "cap must stop at max_delete");
3000 assert!(
3001 r.max_delete_hit,
3002 "max_delete_hit must be true when cap triggered"
3003 );
3004
3005 let mut surviving = 0usize;
3007 for &id in &ids {
3008 if store
3009 .batch_exists(&[id], ns)
3010 .await
3011 .expect("exists")
3012 .contains(&id)
3013 {
3014 surviving += 1;
3015 }
3016 }
3017 assert_eq!(surviving, 3, "3 orphans must survive after cap");
3018 }
3019
3020 #[tokio::test]
3023 async fn orphan_sweep_namespace_filter_scopes_sweep() {
3024 let pool = make_pool();
3025 create_substrate_tables(&pool);
3026 create_vec_table(&pool, "sw_ns", 4);
3027 let store = make_store(Arc::clone(&pool), "sw_ns", 4, "ns:a");
3028
3029 let id_a = Uuid::new_v4();
3030 let id_b = Uuid::new_v4();
3031
3032 store
3033 .insert(
3034 id_a,
3035 SubstrateKind::Entity,
3036 "ns:a",
3037 "body",
3038 vec![vec4(0.1, 0.2, 0.3, 0.4)],
3039 )
3040 .await
3041 .expect("insert ns:a");
3042 store
3043 .insert(
3044 id_b,
3045 SubstrateKind::Entity,
3046 "ns:b",
3047 "body",
3048 vec![vec4(0.5, 0.6, 0.7, 0.8)],
3049 )
3050 .await
3051 .expect("insert ns:b");
3052
3053 let r = store
3055 .orphan_sweep(&OrphanSweepConfig {
3056 subject_id_allowlist: None,
3057 namespaces: vec!["ns:a".to_string()],
3058 substrate_kinds: vec![],
3059 max_delete: 100,
3060 dry_run: false,
3061 })
3062 .await
3063 .expect("sweep");
3064
3065 assert_eq!(r.scanned, 1, "only ns:a row visible to scoped sweep");
3066 assert_eq!(r.deleted, 1);
3067
3068 let exists_a = store.batch_exists(&[id_a], "ns:a").await.expect("exists a");
3069 let exists_b = store.batch_exists(&[id_b], "ns:b").await.expect("exists b");
3070 assert!(!exists_a.contains(&id_a), "ns:a orphan must be swept");
3071 assert!(exists_b.contains(&id_b), "ns:b vec must be untouched");
3072 }
3073
3074 #[tokio::test]
3077 async fn orphan_sweep_substrate_kinds_filter_scopes_sweep() {
3078 let pool = make_pool();
3079 create_substrate_tables(&pool);
3080 create_vec_table(&pool, "sw_kind", 4);
3081 let store = make_store(Arc::clone(&pool), "sw_kind", 4, "ns:kind");
3082 let ns = "ns:kind";
3083
3084 let id_ent = Uuid::new_v4();
3085 let id_note = Uuid::new_v4();
3086
3087 store
3089 .insert(
3090 id_ent,
3091 SubstrateKind::Entity,
3092 ns,
3093 "body",
3094 vec![vec4(0.1, 0.2, 0.3, 0.4)],
3095 )
3096 .await
3097 .expect("insert entity vec");
3098 store
3099 .insert(
3100 id_note,
3101 SubstrateKind::Note,
3102 ns,
3103 "body",
3104 vec![vec4(0.5, 0.6, 0.7, 0.8)],
3105 )
3106 .await
3107 .expect("insert note vec");
3108
3109 let r = store
3111 .orphan_sweep(&OrphanSweepConfig {
3112 subject_id_allowlist: None,
3113 namespaces: vec![],
3114 substrate_kinds: vec![SubstrateKind::Entity],
3115 max_delete: 100,
3116 dry_run: false,
3117 })
3118 .await
3119 .expect("sweep");
3120
3121 assert_eq!(r.scanned, 1, "kind filter restricts scanned count");
3122 assert_eq!(r.deleted, 1, "only entity-kind orphan is swept");
3123
3124 let ent_exists = store.batch_exists(&[id_ent], ns).await.expect("ent exists");
3125 let note_exists = store
3126 .batch_exists(&[id_note], ns)
3127 .await
3128 .expect("note exists");
3129 assert!(
3130 !ent_exists.contains(&id_ent),
3131 "entity-kind orphan must be swept"
3132 );
3133 assert!(
3134 note_exists.contains(&id_note),
3135 "note-kind vec must be untouched"
3136 );
3137 }
3138
3139 #[tokio::test]
3142 async fn orphan_sweep_allowlist_restricts_eligible_rows() {
3143 let pool = make_pool();
3144 create_substrate_tables(&pool);
3145 create_vec_table(&pool, "sw_allow", 4);
3146 let store = make_store(Arc::clone(&pool), "sw_allow", 4, "ns:allow");
3147 let ns = "ns:allow";
3148
3149 let id1 = Uuid::new_v4();
3150 let id2 = Uuid::new_v4();
3151 let id3 = Uuid::new_v4(); for (i, &id) in [id1, id2, id3].iter().enumerate() {
3154 let v = i as f32 * 0.1 + 0.1;
3155 store
3156 .insert(
3157 id,
3158 SubstrateKind::Entity,
3159 ns,
3160 "body",
3161 vec![vec![v, v, v, v]],
3162 )
3163 .await
3164 .expect("insert vec");
3165 }
3166
3167 let r = store
3169 .orphan_sweep(&OrphanSweepConfig {
3170 subject_id_allowlist: Some(vec![id1, id2]),
3171 namespaces: vec![],
3172 substrate_kinds: vec![],
3173 max_delete: 100,
3174 dry_run: false,
3175 })
3176 .await
3177 .expect("sweep");
3178
3179 assert_eq!(r.scanned, 2, "allowlist restricts scanned to 2");
3180 assert_eq!(r.would_delete, 2);
3181 assert_eq!(r.deleted, 2, "both allowlisted orphans deleted");
3182
3183 let e1 = store.batch_exists(&[id1], ns).await.expect("e1");
3184 let e2 = store.batch_exists(&[id2], ns).await.expect("e2");
3185 let e3 = store.batch_exists(&[id3], ns).await.expect("e3");
3186 assert!(!e1.contains(&id1), "id1 must be swept");
3187 assert!(!e2.contains(&id2), "id2 must be swept");
3188 assert!(e3.contains(&id3), "id3 not in allowlist must survive");
3189 }
3190
3191 fn insert_note(pool: &Arc<crate::pool::ConnectionPool>, id: Uuid, deleted_at: Option<i64>) {
3194 let id_str = id.to_string();
3195 pool.try_writer()
3196 .expect("writer")
3197 .conn()
3198 .execute(
3199 "INSERT INTO notes (id, deleted_at) VALUES (?1, ?2)",
3200 rusqlite::params![id_str, deleted_at],
3201 )
3202 .expect("insert note");
3203 }
3204
3205 #[tokio::test]
3208 async fn orphan_sweep_keeps_live_note() {
3209 let pool = make_pool();
3210 create_substrate_tables(&pool);
3211 create_vec_table(&pool, "sw_note_live", 4);
3212 let store = make_store(Arc::clone(&pool), "sw_note_live", 4, "ns:nlive");
3213 let ns = "ns:nlive";
3214
3215 let id = Uuid::new_v4();
3216 insert_note(&pool, id, None); store
3219 .insert(
3220 id,
3221 SubstrateKind::Note,
3222 ns,
3223 "body",
3224 vec![vec4(0.1, 0.2, 0.3, 0.4)],
3225 )
3226 .await
3227 .expect("insert vec");
3228
3229 let r = store
3230 .orphan_sweep(&sweep_all(100, false))
3231 .await
3232 .expect("sweep");
3233
3234 assert_eq!(r.scanned, 1);
3235 assert_eq!(r.would_delete, 0, "live note is not an orphan");
3236 assert_eq!(r.deleted, 0);
3237
3238 let present = store.batch_exists(&[id], ns).await.expect("exists");
3239 assert!(present.contains(&id), "live note's vec must survive");
3240 }
3241
3242 #[tokio::test]
3245 async fn orphan_sweep_sweeps_soft_deleted_note() {
3246 let pool = make_pool();
3247 create_substrate_tables(&pool);
3248 create_vec_table(&pool, "sw_note_soft", 4);
3249 let store = make_store(Arc::clone(&pool), "sw_note_soft", 4, "ns:nsoft");
3250 let ns = "ns:nsoft";
3251
3252 let id = Uuid::new_v4();
3253 insert_note(&pool, id, Some(1_000_000)); store
3256 .insert(
3257 id,
3258 SubstrateKind::Note,
3259 ns,
3260 "body",
3261 vec![vec4(0.5, 0.5, 0.5, 0.5)],
3262 )
3263 .await
3264 .expect("insert vec");
3265
3266 let r = store
3267 .orphan_sweep(&sweep_all(100, false))
3268 .await
3269 .expect("sweep");
3270
3271 assert_eq!(r.scanned, 1);
3272 assert_eq!(r.would_delete, 1, "soft-deleted note counts as orphan");
3273 assert_eq!(r.deleted, 1);
3274
3275 let present = store.batch_exists(&[id], ns).await.expect("exists");
3276 assert!(
3277 !present.contains(&id),
3278 "soft-deleted note's vec must be swept"
3279 );
3280 }
3281
3282 #[tokio::test]
3296 async fn orphan_sweep_error_does_not_poison_connection() {
3297 let pool = make_pool();
3298 create_vec_table(&pool, "sw_poison", 4);
3300 let store = make_store(Arc::clone(&pool), "sw_poison", 4, "ns:poison");
3301 let ns = "ns:poison";
3302
3303 let sweep_result = store.orphan_sweep(&sweep_all(100, false)).await;
3305 assert!(
3306 sweep_result.is_err(),
3307 "sweep must fail when substrate tables are absent"
3308 );
3309
3310 let id = Uuid::new_v4();
3312 store
3313 .insert(
3314 id,
3315 SubstrateKind::Entity,
3316 ns,
3317 "body",
3318 vec![vec4(0.1, 0.2, 0.3, 0.4)],
3319 )
3320 .await
3321 .expect("insert after failed sweep must succeed (connection not poisoned)");
3322
3323 let present = store.batch_exists(&[id], ns).await.expect("exists");
3324 assert!(
3325 present.contains(&id),
3326 "vector inserted after failed sweep must be present"
3327 );
3328 }
3329}
3330
3331#[cfg(all(test, feature = "vectors"))]
3342mod write_queue_tests {
3343 use std::sync::Arc;
3344 use std::time::Duration;
3345
3346 use khive_storage::types::VectorRecord;
3347 use khive_storage::VectorStore;
3348 use khive_types::SubstrateKind;
3349 use uuid::Uuid;
3350
3351 use super::*;
3352 use crate::pool::{ConnectionPool, PoolConfig};
3353
3354 fn create_vec_table(pool: &Arc<ConnectionPool>, model_key: &str, dims: usize) {
3355 let ddl = format!(
3356 "CREATE VIRTUAL TABLE IF NOT EXISTS vec_{} USING vec0(\
3357 subject_id TEXT PRIMARY KEY, \
3358 namespace TEXT NOT NULL, \
3359 kind TEXT NOT NULL, \
3360 field TEXT NOT NULL, \
3361 embedding_model TEXT NOT NULL, \
3362 embedding float[{}] distance_metric=cosine)",
3363 model_key, dims
3364 );
3365 let writer = pool.writer().expect("writer");
3366 writer.conn().execute_batch(&ddl).expect("create vec table");
3367 writer
3368 .conn()
3369 .execute_batch(crate::migrations::ANN_WRITE_LOG_DDL)
3370 .expect("create ann_write_log");
3371 }
3372
3373 #[tokio::test]
3382 async fn insert_batch_routes_through_writer_task_when_flag_enabled() {
3383 crate::extension::ensure_extensions_loaded();
3384
3385 let model_key = "write_queue_flag_test";
3386 let dims = 4usize;
3387 let dir = tempfile::tempdir().unwrap();
3388 let path = dir.path().join("write_queue_vectors.db");
3389 let pool = Arc::new(
3390 ConnectionPool::new(PoolConfig {
3391 path: Some(path),
3392 write_queue_enabled: true,
3393 ..PoolConfig::default()
3394 })
3395 .expect("file-backed pool"),
3396 );
3397 create_vec_table(&pool, model_key, dims);
3398
3399 let store = SqliteVecStore::new(
3400 Arc::clone(&pool),
3401 true,
3402 model_key.to_string(),
3403 model_key.to_string(),
3404 dims,
3405 "ns:test".to_string(),
3406 )
3407 .expect("SqliteVecStore::new");
3408
3409 let id1 = Uuid::new_v4();
3410 let id2 = Uuid::new_v4();
3411 let records = vec![
3412 VectorRecord {
3413 subject_id: id1,
3414 kind: SubstrateKind::Entity,
3415 namespace: "ns:test".to_string(),
3416 field: "body".to_string(),
3417 embedding_model: None,
3418 vectors: vec![vec![0.1, 0.2, 0.3, 0.4]],
3419 updated_at: chrono::Utc::now(),
3420 },
3421 VectorRecord {
3422 subject_id: id2,
3423 kind: SubstrateKind::Entity,
3424 namespace: "ns:test".to_string(),
3425 field: "body".to_string(),
3426 embedding_model: None,
3427 vectors: vec![vec![0.5, 0.6, 0.7, 0.8]],
3428 updated_at: chrono::Utc::now(),
3429 },
3430 ];
3431
3432 let summary = store.insert_batch(records).await.unwrap();
3433 assert_eq!(summary.attempted, 2);
3434 assert_eq!(summary.affected, 2);
3435 assert_eq!(summary.failed, 0);
3436
3437 let present = store
3438 .batch_exists(&[id1, id2], "ns:test")
3439 .await
3440 .expect("batch_exists");
3441 assert!(present.contains(&id1));
3442 assert!(present.contains(&id2));
3443 assert_eq!(
3444 pool.writer_task_spawn_count(),
3445 1,
3446 "the flag-ON path must actually spawn and use the writer task"
3447 );
3448 }
3449
3450 fn create_substrate_tables(pool: &Arc<ConnectionPool>) {
3456 pool.try_writer()
3457 .expect("writer")
3458 .conn()
3459 .execute_batch(
3460 "CREATE TABLE IF NOT EXISTS entities \
3461 (id TEXT PRIMARY KEY, deleted_at INTEGER); \
3462 CREATE TABLE IF NOT EXISTS notes \
3463 (id TEXT PRIMARY KEY, deleted_at INTEGER);",
3464 )
3465 .expect("create substrate tables");
3466 }
3467
3468 fn insert_entity(pool: &Arc<ConnectionPool>, id: Uuid, deleted_at: Option<i64>) {
3470 let id_str = id.to_string();
3471 pool.try_writer()
3472 .expect("writer")
3473 .conn()
3474 .execute(
3475 "INSERT INTO entities (id, deleted_at) VALUES (?1, ?2)",
3476 rusqlite::params![id_str, deleted_at],
3477 )
3478 .expect("insert entity");
3479 }
3480
3481 #[tokio::test]
3486 async fn orphan_sweep_routes_through_writer_task_when_flag_enabled() {
3487 crate::extension::ensure_extensions_loaded();
3488
3489 let model_key = "write_queue_orphan_sweep";
3490 let dims = 4usize;
3491 let dir = tempfile::tempdir().unwrap();
3492 let path = dir.path().join("write_queue_orphan_sweep.db");
3493 let pool = Arc::new(
3494 ConnectionPool::new(PoolConfig {
3495 path: Some(path),
3496 write_queue_enabled: true,
3497 ..PoolConfig::default()
3498 })
3499 .expect("file-backed pool"),
3500 );
3501 create_substrate_tables(&pool);
3502 create_vec_table(&pool, model_key, dims);
3503
3504 let store = SqliteVecStore::new(
3505 Arc::clone(&pool),
3506 true,
3507 model_key.to_string(),
3508 model_key.to_string(),
3509 dims,
3510 "ns:test".to_string(),
3511 )
3512 .expect("SqliteVecStore::new");
3513
3514 let live_id = Uuid::new_v4();
3515 insert_entity(&pool, live_id, None); let orphan_id = Uuid::new_v4(); store
3519 .insert(
3520 live_id,
3521 SubstrateKind::Entity,
3522 "ns:test",
3523 "body",
3524 vec![vec![0.1, 0.2, 0.3, 0.4]],
3525 )
3526 .await
3527 .expect("insert live vector");
3528 store
3529 .insert(
3530 orphan_id,
3531 SubstrateKind::Entity,
3532 "ns:test",
3533 "body",
3534 vec![vec![0.5, 0.6, 0.7, 0.8]],
3535 )
3536 .await
3537 .expect("insert orphan vector");
3538
3539 let dry = store
3541 .orphan_sweep(&OrphanSweepConfig {
3542 subject_id_allowlist: None,
3543 namespaces: vec![],
3544 substrate_kinds: vec![],
3545 max_delete: 100,
3546 dry_run: true,
3547 })
3548 .await
3549 .expect("dry-run sweep");
3550 assert_eq!(dry.scanned, 2);
3551 assert_eq!(dry.would_delete, 1);
3552 assert_eq!(dry.deleted, 0);
3553 assert!(!dry.max_delete_hit);
3554
3555 let real = store
3557 .orphan_sweep(&OrphanSweepConfig {
3558 subject_id_allowlist: None,
3559 namespaces: vec![],
3560 substrate_kinds: vec![],
3561 max_delete: 100,
3562 dry_run: false,
3563 })
3564 .await
3565 .expect("real sweep");
3566 assert_eq!(real.scanned, 2);
3567 assert_eq!(real.would_delete, 1);
3568 assert_eq!(real.deleted, 1);
3569 assert!(!real.max_delete_hit);
3570
3571 let present = store
3572 .batch_exists(&[live_id, orphan_id], "ns:test")
3573 .await
3574 .expect("batch_exists");
3575 assert!(
3576 present.contains(&live_id),
3577 "live vector must survive the sweep"
3578 );
3579 assert!(
3580 !present.contains(&orphan_id),
3581 "orphaned vector must be swept"
3582 );
3583
3584 let writer_task = pool
3603 .writer_task_handle()
3604 .expect("writer task handle")
3605 .expect("writer task must be spawned for a file-backed pool with the flag on");
3606
3607 let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>();
3608 let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
3609 let occupier = {
3610 let writer_task = writer_task.clone();
3611 tokio::spawn(async move {
3612 writer_task
3613 .send(move |_conn| {
3614 let _ = started_tx.send(());
3615 let _ = release_rx.blocking_recv();
3616 Ok::<(), StorageError>(())
3617 })
3618 .await
3619 })
3620 };
3621
3622 started_rx
3623 .await
3624 .expect("occupier must signal it has started running inside the writer task");
3625 assert_eq!(
3626 writer_task.queue_depth(),
3627 0,
3628 "channel must start empty once the occupier has been dequeued and is running"
3629 );
3630
3631 let sweep_task = tokio::spawn(async move {
3632 store
3633 .orphan_sweep(&OrphanSweepConfig {
3634 subject_id_allowlist: None,
3635 namespaces: vec![],
3636 substrate_kinds: vec![],
3637 max_delete: 100,
3638 dry_run: true,
3639 })
3640 .await
3641 });
3642
3643 let mut saw_enqueued = false;
3644 for _ in 0..100 {
3645 if writer_task.queue_depth() >= 1 {
3646 saw_enqueued = true;
3647 break;
3648 }
3649 tokio::time::sleep(Duration::from_millis(5)).await;
3650 }
3651 assert!(
3652 saw_enqueued,
3653 "orphan_sweep's write request never appeared in the writer task's channel \
3654 while the occupier held the single drain slot — orphan_sweep is not routing \
3655 through the shared writer task"
3656 );
3657
3658 release_tx
3659 .send(())
3660 .expect("occupier must still be waiting on the release signal");
3661 occupier
3662 .await
3663 .expect("occupier task must not panic")
3664 .expect("occupier write must succeed");
3665 let post_sweep = sweep_task
3666 .await
3667 .expect("sweep task must not panic")
3668 .expect("orphan_sweep must succeed once unblocked");
3669 assert_eq!(
3670 post_sweep.scanned, 1,
3671 "only the surviving live vector remains after the earlier real sweep"
3672 );
3673 }
3674
3675 #[tokio::test]
3687 async fn orphan_sweep_old_unmanaged_shape_nests_transaction_under_write_queue() {
3688 crate::extension::ensure_extensions_loaded();
3689
3690 let dir = tempfile::tempdir().unwrap();
3691 let path = dir.path().join("write_queue_orphan_sweep_regression.db");
3692 let pool = Arc::new(
3693 ConnectionPool::new(PoolConfig {
3694 path: Some(path),
3695 write_queue_enabled: true,
3696 ..PoolConfig::default()
3697 })
3698 .expect("file-backed pool"),
3699 );
3700 create_substrate_tables(&pool);
3701 create_vec_table(&pool, "write_queue_orphan_sweep_regression", 4);
3702
3703 let writer_task = pool
3704 .writer_task_handle()
3705 .expect("writer task handle")
3706 .expect("writer task must spawn for a file-backed pool with the flag on");
3707
3708 let result: Result<(), StorageError> = writer_task
3709 .send(move |conn| {
3710 let tx = rusqlite::Transaction::new_unchecked(
3716 conn,
3717 rusqlite::TransactionBehavior::Immediate,
3718 )
3719 .map_err(|e| map_err(e, "orphan_sweep_old_shape"))?;
3720 tx.commit()
3721 .map_err(|e| map_err(e, "orphan_sweep_old_shape"))?;
3722 Ok(())
3723 })
3724 .await;
3725
3726 let err = result.expect_err(
3727 "routing the OLD orphan_sweep closure (its own BEGIN IMMEDIATE) through the \
3728 WriterTask must fail under KHIVE_WRITE_QUEUE — if this now succeeds, re-audit \
3729 whether the WriterTask still owns the sole BEGIN IMMEDIATE for this connection",
3730 );
3731 let msg = err.to_string();
3732 assert!(
3733 msg.contains("cannot start a transaction within a transaction"),
3734 "expected the deterministic nested-transaction failure (SQLite's own message \
3735 for a second BEGIN issued inside an already-open transaction), got: {msg}"
3736 );
3737 }
3738}