1use std::collections::HashMap;
41
42use serde_json::Value;
43use uuid::Uuid;
44
45use khive_storage::note::Note;
46use khive_storage::types::SqlValue;
47use khive_storage::{SqlStatement, StorageCapability, StorageError};
48use khive_types::SubstrateKind;
49
50use crate::atomic_plan::{AddNotePlan, AffectedRowGuard, PlanStatement, PostCommitEffect};
51use crate::atomic_runner::{AtomicOpPlan, AtomicRunOutcome};
52use crate::config::NamespaceToken;
53use crate::curation::note_fts_document;
54use crate::error::{RuntimeError, RuntimeResult};
55use crate::runtime::KhiveRuntime;
56
57pub struct AtomicNoteSpec<'a> {
63 pub token: &'a NamespaceToken,
65 pub id: Option<Uuid>,
70 pub kind: &'a str,
71 pub name: Option<&'a str>,
72 pub content: &'a str,
73 pub properties: Option<Value>,
74}
75
76#[derive(Clone, Copy, Default)]
77pub(crate) struct AtomicNoteOptions<'a> {
78 pub salience: Option<f64>,
79 pub decay_factor: Option<f64>,
80 pub embedding_model: Option<&'a str>,
81 pub embedding_content: Option<&'a str>,
82 pub embed: Option<bool>,
83 pub key: Option<&'a str>,
84 pub memory_visibility_receipt: bool,
87 pub replay_receipt: bool,
92 pub fence: Option<&'a crate::note_write::NoteFences>,
93 pub properties_already_derived: bool,
103}
104
105pub(crate) struct AtomicNoteRequest<'a> {
106 pub spec: AtomicNoteSpec<'a>,
107 pub options: AtomicNoteOptions<'a>,
108}
109
110pub(crate) struct PreparedAtomicNotes {
111 pub notes: Vec<Note>,
112 pub plans: Vec<AtomicOpPlan>,
113 pub embedding_truncation: crate::retrieval::EmbeddingTruncationReport,
114}
115
116fn non_finite_index(data: &[f32]) -> Option<usize> {
121 data.iter().position(|v| !v.is_finite())
122}
123
124pub(crate) fn non_finite_vector_error(idx: usize, value: f32) -> RuntimeError {
129 RuntimeError::Storage(StorageError::InvalidInput {
130 capability: StorageCapability::Vectors,
131 operation: "vec_insert".into(),
132 message: format!(
133 "non-finite value at index {idx}: {value} \
134 (NaN/Inf values corrupt distance computations)"
135 ),
136 })
137}
138
139#[cfg(any(test, feature = "fault-injection"))]
152fn injected_failure_statement(label: &str) -> PlanStatement {
153 PlanStatement {
154 statement: SqlStatement {
155 sql: "UPDATE notes SET updated_at = updated_at WHERE 1 = 0".to_string(),
156 params: vec![],
157 label: Some(label.to_string()),
158 },
159 guard: Some(AffectedRowGuard::exactly(1)),
160 }
161}
162
163#[cfg(any(test, feature = "fault-injection"))]
164fn maybe_inject_fts_failure(namespace: &str, label: &str) -> Option<PlanStatement> {
165 crate::operations::consume_fts_fail_fault(namespace).then(|| injected_failure_statement(label))
166}
167#[cfg(not(any(test, feature = "fault-injection")))]
168fn maybe_inject_fts_failure(_namespace: &str, _label: &str) -> Option<PlanStatement> {
169 None
170}
171
172#[cfg(any(test, feature = "fault-injection"))]
173fn maybe_inject_vector_failure(namespace: &str, label: &str) -> Option<PlanStatement> {
174 crate::operations::consume_vector_fail_fault(namespace)
175 .then(|| injected_failure_statement(label))
176}
177#[cfg(not(any(test, feature = "fault-injection")))]
178fn maybe_inject_vector_failure(_namespace: &str, _label: &str) -> Option<PlanStatement> {
179 None
180}
181
182#[allow(clippy::too_many_arguments)]
192pub(crate) fn vector_insert_statements(
193 table: &str,
194 namespace: &str,
195 subject_id: Uuid,
196 field: &str,
197 embedding_model: &str,
198 embedding: &[f32],
199 label_prefix: &str,
200) -> Vec<PlanStatement> {
201 let subject = subject_id.to_string();
202 let model_key = table
203 .strip_prefix("vec_")
204 .expect("runtime vector tables use the vec_ prefix");
205 let kind_str = SubstrateKind::Note.to_string();
206 let blob = khive_storage::encode_f32_native(embedding);
207 vec![
208 PlanStatement {
211 statement: SqlStatement {
212 sql: format!(
213 "INSERT INTO ann_write_log \
214 (namespace, embedding_model, kind, field, subject_id, op) \
215 SELECT namespace, embedding_model, kind, field, subject_id, 'delete' \
216 FROM {table} WHERE subject_id = ?1 AND NOT \
217 (namespace = ?2 AND embedding_model = ?3 AND kind = ?4 AND field = ?5)"
218 ),
219 params: vec![
220 SqlValue::Text(subject.clone()),
221 SqlValue::Text(namespace.to_string()),
222 SqlValue::Text(embedding_model.to_string()),
223 SqlValue::Text(kind_str.clone()),
224 SqlValue::Text(field.to_string()),
225 ],
226 label: Some(format!("{label_prefix}-log-delete")),
227 },
228 guard: None,
229 },
230 PlanStatement {
231 statement: SqlStatement {
232 sql: format!("DELETE FROM {table} WHERE subject_id = ?1"),
233 params: vec![SqlValue::Text(subject.clone())],
234 label: Some(format!("{label_prefix}-delete")),
235 },
236 guard: None,
237 },
238 PlanStatement {
241 statement: SqlStatement {
242 sql: "DELETE FROM vector_provenance \
243 WHERE model_key = ?1 AND subject_id = ?2"
244 .to_string(),
245 params: vec![
246 SqlValue::Text(model_key.to_string()),
247 SqlValue::Text(subject.clone()),
248 ],
249 label: Some(format!("{label_prefix}-clear-provenance")),
250 },
251 guard: None,
252 },
253 PlanStatement {
254 statement: SqlStatement {
255 sql: format!(
256 "INSERT INTO {table} \
257 (subject_id, namespace, kind, field, embedding_model, embedding) \
258 VALUES (?1, ?2, ?3, ?4, ?5, ?6)"
259 ),
260 params: vec![
261 SqlValue::Text(subject.clone()),
262 SqlValue::Text(namespace.to_string()),
263 SqlValue::Text(kind_str.clone()),
264 SqlValue::Text(field.to_string()),
265 SqlValue::Text(embedding_model.to_string()),
266 SqlValue::Blob(blob),
267 ],
268 label: Some(format!("{label_prefix}-insert")),
269 },
270 guard: None,
271 },
272 PlanStatement {
273 statement: SqlStatement {
274 sql: "INSERT INTO ann_write_log \
275 (namespace, embedding_model, kind, field, subject_id, op) \
276 VALUES (?1, ?2, ?3, ?4, ?5, 'upsert')"
277 .to_string(),
278 params: vec![
279 SqlValue::Text(namespace.to_string()),
280 SqlValue::Text(embedding_model.to_string()),
281 SqlValue::Text(kind_str),
282 SqlValue::Text(field.to_string()),
283 SqlValue::Text(subject),
284 ],
285 label: Some(format!("{label_prefix}-log-upsert")),
286 },
287 guard: None,
288 },
289 ]
290}
291
292#[cfg(test)]
299pub(crate) async fn create_notes_atomic(
300 runtime: &KhiveRuntime,
301 specs: Vec<AtomicNoteSpec<'_>>,
302) -> RuntimeResult<Vec<Note>> {
303 Ok(create_notes_atomic_with_report(runtime, specs).await?.0)
304}
305
306pub async fn create_notes_atomic_with_report(
311 runtime: &KhiveRuntime,
312 specs: Vec<AtomicNoteSpec<'_>>,
313) -> RuntimeResult<(Vec<Note>, crate::retrieval::EmbeddingTruncationReport)> {
314 create_notes_atomic_with_attachments(runtime, specs, &[]).await
315}
316
317pub async fn create_notes_atomic_with_attachments(
320 runtime: &KhiveRuntime,
321 specs: Vec<AtomicNoteSpec<'_>>,
322 attachments: &[khive_storage::NewAttachment],
323) -> RuntimeResult<(Vec<Note>, crate::retrieval::EmbeddingTruncationReport)> {
324 validate_note_attachments(runtime, attachments)?;
325 let mut prepared = prepare_atomic_notes(runtime, specs, AtomicNoteOptions::default()).await?;
326 append_note_attachments(&mut prepared, attachments)?;
327 match crate::atomic_runner::run_atomic_unit_with_note_versions(
328 runtime.sql().as_ref(),
329 prepared.plans,
330 true,
331 )
332 .await
333 {
334 Ok((AtomicRunOutcome::Committed { .. }, versions)) => {
335 assert_eq!(
336 prepared.notes.len(),
337 versions.len(),
338 "one revision receipt per prepared note"
339 );
340 for (note, version) in prepared.notes.iter_mut().zip(versions) {
341 note.version = version;
342 }
343 Ok((prepared.notes, prepared.embedding_truncation))
344 }
345 Ok((
346 AtomicRunOutcome::RolledBack {
347 failed_op_index,
348 failure,
349 },
350 _,
351 )) => Err(RuntimeError::Internal(format!(
352 "atomic multi-note write rolled back at op {failed_op_index}: {failure:?}"
353 ))),
354 Err(e) => Err(RuntimeError::Storage(e.0)),
355 }
356}
357
358pub(crate) fn validate_note_attachments(
359 runtime: &KhiveRuntime,
360 attachments: &[khive_storage::NewAttachment],
361) -> RuntimeResult<()> {
362 if !attachments.is_empty() {
363 runtime.attachments()?;
366 }
367 for attachment in attachments {
368 attachment.validate()?;
369 }
370 Ok(())
371}
372
373pub(crate) fn append_note_attachments(
374 prepared: &mut PreparedAtomicNotes,
375 attachments: &[khive_storage::NewAttachment],
376) -> RuntimeResult<()> {
377 for (plan, note) in prepared.plans.iter_mut().zip(&prepared.notes) {
378 let AtomicOpPlan::AddNote(plan) = plan else {
379 return Err(RuntimeError::Internal(
380 "expected prepared attachment owner note".into(),
381 ));
382 };
383 for attachment in attachments {
384 let row = khive_storage::Attachment::from_new(
385 note.id,
386 khive_storage::AttachmentSubstrate::Note,
387 attachment.clone(),
388 note.created_at,
389 );
390 plan.statements.push(PlanStatement {
393 statement: khive_db::stores::attachment::attachment_upsert_statement(&row)?,
394 guard: Some(AffectedRowGuard::exactly(1)),
395 });
396 }
397 }
398 Ok(())
399}
400
401pub(crate) async fn prepare_atomic_notes(
402 runtime: &KhiveRuntime,
403 specs: Vec<AtomicNoteSpec<'_>>,
404 options: AtomicNoteOptions<'_>,
405) -> RuntimeResult<PreparedAtomicNotes> {
406 validate_atomic_note_options(&options)?;
407 if options.embed != Some(false) {
408 if let Some(model) = options.embedding_model {
409 runtime.resolve_embedding_model(Some(model))?;
410 }
411 }
412 prepare_atomic_note_requests(
413 runtime,
414 specs
415 .into_iter()
416 .map(|spec| AtomicNoteRequest { spec, options })
417 .collect(),
418 )
419 .await
420}
421
422fn validate_atomic_note_options(options: &AtomicNoteOptions<'_>) -> RuntimeResult<()> {
423 if let Some(value) = options.salience {
424 if !value.is_finite() || !(0.0..=1.0).contains(&value) {
425 return Err(RuntimeError::InvalidInput(
426 "salience must be a finite value in [0.0, 1.0]".into(),
427 ));
428 }
429 }
430 if let Some(value) = options.decay_factor {
431 if !value.is_finite() || value < 0.0 {
432 return Err(RuntimeError::InvalidInput(
433 "decay_factor must be a finite value >= 0.0".into(),
434 ));
435 }
436 }
437 Ok(())
438}
439
440pub(crate) async fn prepare_atomic_note_requests(
441 runtime: &KhiveRuntime,
442 requests: Vec<AtomicNoteRequest<'_>>,
443) -> RuntimeResult<PreparedAtomicNotes> {
444 let mut embed_model_names = Vec::new();
445 let mut note_models = Vec::with_capacity(requests.len());
446 for request in &requests {
447 let options = &request.options;
448 validate_atomic_note_options(options)?;
449 let models = if options.embed == Some(false) {
450 Vec::new()
451 } else if let Some(model) = options.embedding_model {
452 runtime.resolve_embedding_model(Some(model))?;
453 vec![model.to_owned()]
454 } else {
455 runtime.embedding_models_for_note_kind(request.spec.kind)
456 };
457 let mut indices = Vec::with_capacity(models.len());
458 for model in models {
459 let index = match embed_model_names.iter().position(|name| name == &model) {
460 Some(index) => index,
461 None => {
462 embed_model_names.push(model);
463 embed_model_names.len() - 1
464 }
465 };
466 indices.push(index);
467 }
468 note_models.push(indices);
469 }
470 let mut notes: Vec<Note> = Vec::with_capacity(requests.len());
473 for (index, request) in requests.iter().enumerate() {
474 let spec = &request.spec;
475 let options = &request.options;
476 let record = format!("note[{index}]");
477 runtime.validate_note_kind(spec.kind)?;
478 let properties = if options.properties_already_derived {
487 spec.properties.clone()
488 } else {
489 runtime.derive_note_write_properties(spec.kind, spec.token, spec.properties.clone())?
490 };
491 crate::secret_gate::reject_reserved_secret_gate_property(properties.as_ref())?;
492 crate::secret_gate::check_at(spec.content, &record, "content")?;
493 if let Some(n) = spec.name {
494 crate::secret_gate::check_at(n, &record, "name")?;
495 }
496 if let Some(ref p) = properties {
497 crate::secret_gate::check_json_at(p, &record, "properties")?;
498 }
499
500 let ns = spec.token.namespace().as_str();
501 let mut note = Note::new(ns, spec.kind, spec.content);
502 note.key = options.key.map(str::to_owned);
503 if let Some(salience) = options.salience {
504 note = note.with_salience(salience);
505 }
506 if let Some(decay_factor) = options.decay_factor {
507 note = note.with_decay(decay_factor);
508 }
509 if let Some(id) = spec.id {
510 note.id = id;
511 }
512 if let Some(n) = spec.name {
513 note = note.with_name(n);
514 }
515 if let Some(p) = properties {
516 note = note.with_properties(p);
517 }
518 notes.push(note);
519 }
520
521 let memory_receipts: Vec<bool> = notes
525 .iter()
526 .zip(&requests)
527 .map(|(note, request)| {
528 note.kind == "memory"
529 && (request.options.key.is_some() || request.options.memory_visibility_receipt)
530 })
531 .collect();
532 if memory_receipts.iter().any(|capture| *capture) {
533 runtime.require_visibility_cutover()?;
534 }
535
536 let mut content_group_by_text: HashMap<&str, usize> = HashMap::new();
540 let mut content_groups: Vec<Vec<usize>> = Vec::new();
541 let mut note_content_groups: Vec<usize> = Vec::with_capacity(notes.len());
542 for (note_idx, note) in notes.iter().enumerate() {
543 let text = requests[note_idx]
544 .options
545 .embedding_content
546 .unwrap_or_else(|| crate::curation::note_embedding_text_ref(note));
547 let content_group_idx = match content_group_by_text.get(text) {
548 Some(&idx) => {
549 content_groups[idx].push(note_idx);
550 idx
551 }
552 None => {
553 let idx = content_groups.len();
554 content_group_by_text.insert(text, idx);
555 content_groups.push(vec![note_idx]);
556 idx
557 }
558 };
559 note_content_groups.push(content_group_idx);
560 }
561
562 let mut embedding_outcomes: Vec<Vec<Option<crate::retrieval::DocumentEmbeddingOutcome>>> =
564 content_groups
565 .iter()
566 .map(|_| vec![None; embed_model_names.len()])
567 .collect();
568 let mut embedding_truncation = crate::retrieval::EmbeddingTruncationReport::default();
569
570 if !embed_model_names.is_empty() {
575 let models = embed_model_names
579 .iter()
580 .map(|name| {
581 runtime
582 .vector_model_metadata(name)
583 .map(|(name, dimensions)| (crate::config::sanitize_key(&name), dimensions))
584 })
585 .collect::<RuntimeResult<Vec<_>>>()?;
586 let model_specs: Vec<_> = models
587 .iter()
588 .map(|(key, dimensions)| (key.as_str(), *dimensions))
589 .collect();
590 runtime.backend().ensure_vector_tables(&model_specs)?;
591
592 let usage_ctx = crate::usage::current();
593 let mut join_set = tokio::task::JoinSet::new();
594 for (model_idx, model_name) in embed_model_names.iter().enumerate() {
595 let mut groups = Vec::new();
596 let mut texts = Vec::new();
597 let mut token = None;
598 for (group_idx, note_indices) in content_groups.iter().enumerate() {
599 let Some(¬e_idx) = note_indices
600 .iter()
601 .find(|&&index| note_models[index].contains(&model_idx))
602 else {
603 continue;
604 };
605 token.get_or_insert_with(|| requests[note_idx].spec.token.clone());
606 groups.push(group_idx);
607 texts.push(
608 requests[note_idx]
609 .options
610 .embedding_content
611 .unwrap_or_else(|| {
612 crate::curation::note_embedding_text_ref(¬es[note_idx])
613 })
614 .to_owned(),
615 );
616 }
617 let rt = runtime.clone();
618 let token = token.expect("a selected model has at least one note");
619 let name = model_name.clone();
620 let ctx = usage_ctx.clone();
621 join_set.spawn(crate::runtime::inherit_request_embedder_scope(async move {
622 let fut = async {
623 let mut outcomes = Vec::with_capacity(texts.len());
624 for chunk in texts.chunks(lattice_embed::DEFAULT_MAX_BATCH_SIZE) {
625 outcomes.extend(
626 rt.embed_document_batch_with_model_outcomes_for_token(
627 &token, &name, chunk,
628 )
629 .await?,
630 );
631 }
632 Ok::<_, RuntimeError>(outcomes)
633 };
634 let result = match ctx {
635 Some(ctx) => crate::usage::scope(ctx, fut).await,
636 None => fut.await,
637 };
638 (groups, model_idx, result)
639 }));
640 }
641
642 while let Some(joined) = join_set.join_next().await {
643 match joined {
644 Ok((groups, model_idx, Ok(outcomes))) => {
645 for (group_idx, outcome) in groups.into_iter().zip(outcomes) {
646 embedding_outcomes[group_idx][model_idx] = Some(outcome);
647 }
648 }
649 Ok((_, _, Err(e))) => {
650 join_set.abort_all();
651 return Err(e);
652 }
653 Err(join_err) => {
654 join_set.abort_all();
655 return Err(RuntimeError::Internal(format!(
656 "embed task panicked: {join_err}"
657 )));
658 }
659 }
660 }
661 }
662
663 for (note_idx, &content_group_idx) in note_content_groups.iter().enumerate() {
666 for &model_idx in ¬e_models[note_idx] {
667 let outcome = embedding_outcomes[content_group_idx][model_idx]
668 .as_ref()
669 .expect("every requested embedding was observed");
670 embedding_truncation.observe(outcome);
671 if outcome.truncated {
672 tracing::warn!(
673 model = %outcome.model_name,
674 source_bytes = outcome.source_bytes,
675 embedded_bytes = outcome.embedded_bytes,
676 "atomic note embedding input truncated; full content will be stored unchanged"
677 );
678 }
679 }
680 }
681
682 for outcomes_for_content in &embedding_outcomes {
688 for outcome in outcomes_for_content.iter().flatten() {
689 if let Some(idx) = non_finite_index(&outcome.vector) {
690 return Err(non_finite_vector_error(idx, outcome.vector[idx]));
691 }
692 }
693 }
694
695 let mut plans: Vec<AtomicOpPlan> = Vec::with_capacity(notes.len());
698 for (note_idx, note) in notes.iter().enumerate() {
699 let outcomes_for_note = &embedding_outcomes[note_content_groups[note_idx]];
700 let mut statements = vec![PlanStatement {
701 statement: if note.key.is_some() {
702 khive_db::stores::note::note_insert_keyed_statement(note)
703 } else {
704 khive_db::stores::note::note_upsert_statement(note)
705 },
706 guard: Some(AffectedRowGuard::exactly(1)),
707 }];
708
709 if memory_receipts[note_idx] {
710 statements.push(PlanStatement {
711 statement: SqlStatement {
712 sql:
713 "INSERT INTO memory_visibility_receipts (namespace, note_id, model_count) \
714 VALUES (?1, ?2, ?3)"
715 .into(),
716 params: vec![
717 SqlValue::Text(note.namespace.clone()),
718 SqlValue::Text(note.id.to_string()),
719 SqlValue::Integer(note_models[note_idx].len() as i64),
720 ],
721 label: Some("memory-visibility-receipt".into()),
722 },
723 guard: Some(AffectedRowGuard::exactly(1)),
724 });
725 }
726
727 if let Some(fault) = maybe_inject_fts_failure(¬e.namespace, "fault-injected-fts") {
728 statements.push(fault);
729 } else {
730 statements.push(PlanStatement {
740 statement: khive_db::stores::text::delete_document_statement(
741 "fts_notes",
742 ¬e.namespace,
743 note.id,
744 ),
745 guard: None,
746 });
747 for statement in khive_db::stores::text::insert_document_statements(
753 "fts_notes",
754 ¬e_fts_document(note),
755 ) {
756 statements.push(PlanStatement {
757 statement,
758 guard: None,
759 });
760 }
761 }
762
763 if let Some(fault) = maybe_inject_vector_failure(¬e.namespace, "fault-injected-vector") {
764 statements.push(fault);
765 } else {
766 for &model_idx in ¬e_models[note_idx] {
767 let model_name = &embed_model_names[model_idx];
768 let outcome = outcomes_for_note[model_idx]
769 .as_ref()
770 .expect("every model index observed exactly once");
771 let table = format!("vec_{}", crate::config::sanitize_key(model_name));
772 statements.extend(vector_insert_statements(
773 &table,
774 ¬e.namespace,
775 note.id,
776 "note.content",
777 model_name,
778 &outcome.vector,
779 &format!("atomic-message-vec-{table}-{}", note.id),
780 ));
781 if memory_receipts[note_idx] {
782 statements.push(PlanStatement {
787 statement: SqlStatement {
788 sql: "INSERT INTO memory_visibility_fences \
789 (namespace, note_id, model, ann_write_log_seq) \
790 VALUES (?1, ?2, ?3, last_insert_rowid())"
791 .into(),
792 params: vec![
793 SqlValue::Text(note.namespace.clone()),
794 SqlValue::Text(note.id.to_string()),
795 SqlValue::Text(model_name.clone()),
796 ],
797 label: Some("memory-visibility-fence".into()),
798 },
799 guard: Some(AffectedRowGuard::exactly(1)),
800 });
801 }
802 }
803 }
804
805 if memory_receipts[note_idx] {
808 statements.push(PlanStatement {
809 statement: SqlStatement {
810 sql: "INSERT INTO memory_visibility_epochs (note_id, namespace, epoch) \
811 VALUES (?1, ?2, 'modern')"
812 .into(),
813 params: vec![
814 SqlValue::Text(note.id.to_string()),
815 SqlValue::Text(note.namespace.clone()),
816 ],
817 label: Some("memory-visibility-epoch".into()),
818 },
819 guard: Some(AffectedRowGuard::exactly(1)),
820 });
821 }
822
823 plans.push(AtomicOpPlan::AddNote(Box::new(AddNotePlan {
824 note_guard: Some(crate::note_write::NoteWriteGuard {
825 namespace: requests[note_idx].spec.token.namespace().as_str().into(),
826 target_id: note.id,
827 expected_version: None,
828 fence: requests[note_idx].options.fence.cloned(),
829 create_key: note
830 .key
831 .as_ref()
832 .map(|key| crate::note_write::CreateKeyClaim {
833 kind: note.kind.clone(),
834 key: key.clone(),
835 content: note.content.clone(),
836 properties: note.properties.clone(),
837 replay_signal: requests[note_idx].options.replay_receipt,
838 }),
839 }),
840 note_id: note.id,
841 statements,
842 post_commit: PostCommitEffect::None,
843 })));
844 }
845
846 Ok(PreparedAtomicNotes {
847 notes,
848 plans,
849 embedding_truncation,
850 })
851}
852
853#[cfg(test)]
854mod tests {
855 use super::*;
856 use async_trait::async_trait;
857 use lattice_embed::{EmbedError, EmbeddingModel, EmbeddingService, MAX_TEXT_BYTES};
858 use std::sync::{Arc, Mutex};
859
860 use khive_types::Namespace;
861
862 use crate::embedder_registry::EmbedderProvider;
863
864 const NAN_MODEL: &str = "atomic-message-nan-model";
865 const NAN_DIMS: usize = 4;
866
867 struct NanService;
868 #[async_trait]
869 impl EmbeddingService for NanService {
870 async fn embed(
871 &self,
872 texts: &[String],
873 _model: EmbeddingModel,
874 ) -> Result<Vec<Vec<f32>>, EmbedError> {
875 Ok(texts.iter().map(|_| vec![f32::NAN; NAN_DIMS]).collect())
876 }
877 fn supports_model(&self, _model: EmbeddingModel) -> bool {
878 true
879 }
880 fn name(&self) -> &'static str {
881 NAN_MODEL
882 }
883 }
884 struct NanProvider;
885 #[async_trait]
886 impl EmbedderProvider for NanProvider {
887 fn name(&self) -> &str {
888 NAN_MODEL
889 }
890 fn dimensions(&self) -> usize {
891 NAN_DIMS
892 }
893 async fn build(&self) -> RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
894 Ok(std::sync::Arc::new(NanService))
895 }
896 }
897
898 struct TruncationService;
899 #[async_trait]
900 impl EmbeddingService for TruncationService {
901 async fn embed(
902 &self,
903 texts: &[String],
904 _model: EmbeddingModel,
905 ) -> Result<Vec<Vec<f32>>, EmbedError> {
906 Ok(texts
907 .iter()
908 .map(|_| vec![0.5; EmbeddingModel::MultilingualE5Base.dimensions()])
909 .collect())
910 }
911 fn supports_model(&self, _model: EmbeddingModel) -> bool {
912 true
913 }
914 fn name(&self) -> &'static str {
915 "multilingual-e5-base"
916 }
917 }
918
919 struct TruncationProvider;
920 #[async_trait]
921 impl EmbedderProvider for TruncationProvider {
922 fn name(&self) -> &str {
923 "multilingual-e5-base"
924 }
925 fn dimensions(&self) -> usize {
926 EmbeddingModel::MultilingualE5Base.dimensions()
927 }
928 async fn build(&self) -> RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
929 Ok(std::sync::Arc::new(TruncationService))
930 }
931 }
932
933 const DEDUP_DIMS: usize = 4;
934
935 struct DedupService {
936 name: &'static str,
937 }
938
939 #[async_trait]
940 impl EmbeddingService for DedupService {
941 async fn embed(
942 &self,
943 texts: &[String],
944 _model: EmbeddingModel,
945 ) -> Result<Vec<Vec<f32>>, EmbedError> {
946 Ok(texts.iter().map(|_| vec![0.25; DEDUP_DIMS]).collect())
947 }
948
949 fn supports_model(&self, _model: EmbeddingModel) -> bool {
950 true
951 }
952
953 fn name(&self) -> &'static str {
954 self.name
955 }
956 }
957
958 struct DedupProvider {
959 name: &'static str,
960 }
961
962 #[async_trait]
963 impl EmbedderProvider for DedupProvider {
964 fn name(&self) -> &str {
965 self.name
966 }
967
968 fn dimensions(&self) -> usize {
969 DEDUP_DIMS
970 }
971
972 async fn build(&self) -> RuntimeResult<std::sync::Arc<dyn EmbeddingService>> {
973 Ok(std::sync::Arc::new(DedupService { name: self.name }))
974 }
975 }
976
977 struct BatchCountingService {
978 name: &'static str,
979 calls: Arc<Mutex<Vec<Vec<String>>>>,
980 }
981
982 #[async_trait]
983 impl EmbeddingService for BatchCountingService {
984 async fn embed(
985 &self,
986 texts: &[String],
987 _model: EmbeddingModel,
988 ) -> Result<Vec<Vec<f32>>, EmbedError> {
989 self.calls.lock().unwrap().push(texts.to_vec());
990 Ok(texts.iter().map(|_| vec![0.25; DEDUP_DIMS]).collect())
991 }
992
993 fn supports_model(&self, _model: EmbeddingModel) -> bool {
994 true
995 }
996
997 fn name(&self) -> &'static str {
998 self.name
999 }
1000 }
1001
1002 struct BatchCountingProvider {
1003 name: &'static str,
1004 calls: Arc<Mutex<Vec<Vec<String>>>>,
1005 }
1006
1007 #[async_trait]
1008 impl EmbedderProvider for BatchCountingProvider {
1009 fn name(&self) -> &str {
1010 self.name
1011 }
1012
1013 fn dimensions(&self) -> usize {
1014 DEDUP_DIMS
1015 }
1016
1017 async fn build(&self) -> RuntimeResult<Arc<dyn EmbeddingService>> {
1018 Ok(Arc::new(BatchCountingService {
1019 name: self.name,
1020 calls: Arc::clone(&self.calls),
1021 }))
1022 }
1023 }
1024
1025 #[tokio::test]
1026 async fn heterogeneous_notes_batch_models_with_one_schema_writer() {
1027 let runtime = KhiveRuntime::memory().unwrap();
1028 let token = runtime
1029 .authorize(Namespace::parse("batch-notes").unwrap())
1030 .unwrap();
1031 let names = ["batch-notes-a", "batch-notes-b"];
1032 let mut counters = Vec::new();
1033 for name in names {
1034 let calls = Arc::new(Mutex::new(Vec::new()));
1035 runtime.register_embedder(BatchCountingProvider {
1036 name,
1037 calls: Arc::clone(&calls),
1038 });
1039 runtime.embedder_with_token(&token, name).await.unwrap();
1040 counters.push(calls);
1041 }
1042 let inputs = [
1043 ("first", "shared text", true, 0.2),
1044 ("second", "shared text", true, 0.4),
1045 ("third", "different text", true, 0.6),
1046 ("disabled", "shared text", false, 0.8),
1047 ];
1048 let requests = inputs
1049 .iter()
1050 .map(|&(key, content, embed, salience)| AtomicNoteRequest {
1051 spec: AtomicNoteSpec {
1052 token: &token,
1053 id: None,
1054 kind: "observation",
1055 name: None,
1056 content,
1057 properties: Some(serde_json::json!({"key": key})),
1058 },
1059 options: AtomicNoteOptions {
1060 key: Some(key),
1061 salience: Some(salience),
1062 embed: Some(embed),
1063 ..Default::default()
1064 },
1065 })
1066 .collect();
1067 let before = runtime.backend().pool().writer_acquisition_snapshot();
1068 let prepared = prepare_atomic_note_requests(&runtime, requests)
1069 .await
1070 .unwrap();
1071 let after = runtime.backend().pool().writer_acquisition_snapshot();
1072 assert_eq!(after.acquisitions - before.acquisitions, 1);
1073 for calls in &counters {
1074 assert_eq!(
1075 *calls.lock().unwrap(),
1076 vec![vec!["shared text".to_owned(), "different text".to_owned()]],
1077 );
1078 }
1079 for (note, &(key, content, _, salience)) in prepared.notes.iter().zip(&inputs) {
1080 assert_eq!(note.key.as_deref(), Some(key));
1081 assert_eq!(note.content, content);
1082 assert_eq!(note.salience, Some(salience));
1083 assert_eq!(note.properties.as_ref().unwrap()["key"], key);
1084 }
1085 assert!(matches!(
1086 crate::atomic_runner::run_atomic_unit(runtime.sql().as_ref(), prepared.plans)
1087 .await
1088 .unwrap(),
1089 AtomicRunOutcome::Committed { .. }
1090 ));
1091 for name in names {
1092 assert_eq!(
1093 runtime
1094 .vectors_for_model(&token, name)
1095 .unwrap()
1096 .count()
1097 .await
1098 .unwrap(),
1099 3,
1100 "disabled notes must not inherit a sibling's shared vector",
1101 );
1102 }
1103 }
1104
1105 #[tokio::test]
1106 async fn atomic_note_batch_chunks_only_above_provider_limit() {
1107 let runtime = KhiveRuntime::memory().unwrap();
1108 let token = runtime.authorize(Namespace::local()).unwrap();
1109 let calls = Arc::new(Mutex::new(Vec::new()));
1110 runtime.register_embedder(BatchCountingProvider {
1111 name: "batch-notes-large",
1112 calls: Arc::clone(&calls),
1113 });
1114 let contents: Vec<_> = (0..=lattice_embed::DEFAULT_MAX_BATCH_SIZE)
1115 .map(|index| format!("distinct document {index}"))
1116 .collect();
1117 let specs = contents
1118 .iter()
1119 .map(|content| AtomicNoteSpec {
1120 token: &token,
1121 id: None,
1122 kind: "observation",
1123 name: None,
1124 content,
1125 properties: None,
1126 })
1127 .collect();
1128 let prepared = prepare_atomic_notes(&runtime, specs, AtomicNoteOptions::default())
1129 .await
1130 .unwrap();
1131 assert_eq!(prepared.notes.len(), contents.len());
1132 let calls = calls.lock().unwrap();
1133 assert_eq!(calls.len(), 2);
1134 assert_eq!(calls[0].len(), lattice_embed::DEFAULT_MAX_BATCH_SIZE);
1135 assert_eq!(calls[1].len(), 1);
1136 assert_eq!(calls.concat(), contents);
1137 }
1138
1139 fn stamp_from_actor(
1144 _kind: &str,
1145 actor_id: &str,
1146 properties: Option<Value>,
1147 ) -> RuntimeResult<Option<Value>> {
1148 let mut props = match properties {
1149 Some(Value::Object(map)) => map,
1150 _ => serde_json::Map::new(),
1151 };
1152 props.insert(
1153 "from_actor".to_string(),
1154 Value::String(actor_id.to_string()),
1155 );
1156 Ok(Some(Value::Object(props)))
1157 }
1158
1159 async fn fts_row_count(runtime: &KhiveRuntime, namespace: &str) -> i64 {
1160 let mut reader = runtime.sql().reader().await.expect("sql reader");
1161 match reader
1162 .query_scalar(SqlStatement {
1163 sql: "SELECT COUNT(*) FROM fts_notes WHERE namespace = ?1".to_string(),
1164 params: vec![SqlValue::Text(namespace.to_string())],
1165 label: None,
1166 })
1167 .await
1168 .expect("fts count query")
1169 {
1170 Some(SqlValue::Integer(n)) => n,
1171 other => panic!("unexpected fts count result: {other:?}"),
1172 }
1173 }
1174
1175 async fn ann_write_log_count(runtime: &KhiveRuntime, namespace: &str, model: &str) -> i64 {
1176 let mut reader = runtime.sql().reader().await.expect("sql reader");
1177 match reader
1178 .query_scalar(SqlStatement {
1179 sql: "SELECT COUNT(*) FROM ann_write_log \
1180 WHERE namespace = ?1 AND embedding_model = ?2"
1181 .to_string(),
1182 params: vec![
1183 SqlValue::Text(namespace.to_string()),
1184 SqlValue::Text(model.to_string()),
1185 ],
1186 label: None,
1187 })
1188 .await
1189 .expect("ann_write_log count query")
1190 {
1191 Some(SqlValue::Integer(n)) => n,
1192 other => panic!("unexpected ann_write_log count result: {other:?}"),
1193 }
1194 }
1195
1196 #[tokio::test]
1201 async fn create_notes_atomic_rejects_non_finite_embedding_before_any_write() {
1202 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
1203 runtime.register_embedder(NanProvider);
1204 let ns = "atomic-message-nan-test";
1205 let token = runtime
1206 .authorize(Namespace::parse(ns).unwrap())
1207 .expect("authorize");
1208
1209 let result = create_notes_atomic(
1210 &runtime,
1211 vec![AtomicNoteSpec {
1212 token: &token,
1213 id: None,
1214 kind: "observation",
1215 name: None,
1216 content: "nan embedding content",
1217 properties: None,
1218 }],
1219 )
1220 .await;
1221
1222 assert!(
1223 result.is_err(),
1224 "a non-finite embedding must be rejected; got {result:?}"
1225 );
1226
1227 let alive = runtime
1228 .list_notes(&token, Some("observation"), 100, 0)
1229 .await
1230 .expect("list_notes")
1231 .into_iter()
1232 .filter(|n| n.deleted_at.is_none())
1233 .count();
1234 assert_eq!(
1235 alive, 0,
1236 "no note row may be committed when an embedding is non-finite"
1237 );
1238
1239 assert_eq!(
1240 fts_row_count(&runtime, ns).await,
1241 0,
1242 "no FTS document may be committed when an embedding is non-finite"
1243 );
1244
1245 let vs = runtime
1246 .vectors_for_model(&token, NAN_MODEL)
1247 .expect("vec store");
1248 assert_eq!(
1249 vs.count().await.expect("count"),
1250 0,
1251 "no vector row may be committed when an embedding is non-finite"
1252 );
1253
1254 assert_eq!(
1255 ann_write_log_count(&runtime, ns, NAN_MODEL).await,
1256 0,
1257 "no ann_write_log row may be committed when an embedding is non-finite"
1258 );
1259 }
1260
1261 #[tokio::test]
1262 async fn create_notes_atomic_rejects_reserved_secret_gate_key() {
1263 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
1264 let ns = "atomic-message-reserved-key-test";
1265 let token = runtime
1266 .authorize(Namespace::parse(ns).unwrap())
1267 .expect("authorize");
1268
1269 let result = create_notes_atomic(
1270 &runtime,
1271 vec![AtomicNoteSpec {
1272 token: &token,
1273 id: None,
1274 kind: "observation",
1275 name: None,
1276 content: "atomic message reservation target",
1277 properties: Some(
1278 serde_json::json!({"khive:secret_gate": "exempted:content-sha256-manifest-v1"}),
1279 ),
1280 }],
1281 )
1282 .await;
1283
1284 let err = result.expect_err("caller-supplied reserved key must be rejected");
1285 assert!(
1286 matches!(err, RuntimeError::InvalidInput(ref msg) if msg.contains("khive:secret_gate")),
1287 "unexpected error: {err:?}"
1288 );
1289
1290 let alive = runtime
1291 .list_notes(&token, Some("observation"), 100, 0)
1292 .await
1293 .expect("list_notes")
1294 .into_iter()
1295 .filter(|n| n.deleted_at.is_none())
1296 .count();
1297 assert_eq!(
1298 alive, 0,
1299 "no note row may be committed on reservation rejection"
1300 );
1301 }
1302
1303 #[tokio::test]
1304 async fn create_notes_atomic_with_report_preserves_truncation_outcome() {
1305 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
1306 runtime.register_embedder(TruncationProvider);
1307 let token = runtime
1308 .authorize(Namespace::parse("atomic-message-truncation-test").unwrap())
1309 .expect("authorize");
1310 let content = "x".repeat(MAX_TEXT_BYTES);
1311
1312 let (notes, report) = create_notes_atomic_with_report(
1313 &runtime,
1314 vec![AtomicNoteSpec {
1315 token: &token,
1316 id: None,
1317 kind: "observation",
1318 name: None,
1319 content: &content,
1320 properties: None,
1321 }],
1322 )
1323 .await
1324 .expect("atomic note write");
1325
1326 assert_eq!(notes.len(), 1);
1327 assert_eq!(report.truncated, 1);
1328 assert_eq!(
1329 report.discarded_bytes,
1330 "passage: ".len() as u64,
1331 "E5 document-prefix reservation must be reflected in the returned report"
1332 );
1333 }
1334
1335 #[tokio::test]
1336 async fn create_notes_atomic_reuses_identical_content_once_per_model() {
1337 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
1338 for name in ["atomic-dedup-model-a", "atomic-dedup-model-b"] {
1339 runtime.register_embedder(DedupProvider { name });
1340 }
1341 let outbound_token = runtime
1342 .authorize(Namespace::parse("atomic-dedup-outbound").unwrap())
1343 .expect("authorize outbound");
1344 let inbound_token = runtime
1345 .authorize(Namespace::parse("atomic-dedup-inbound").unwrap())
1346 .expect("authorize inbound");
1347 let content = "byte-identical outbound and inbound message content";
1348
1349 let usage = crate::usage::UsageContext::new();
1350 let notes = crate::usage::scope(usage.clone(), async {
1351 create_notes_atomic(
1352 &runtime,
1353 vec![
1354 AtomicNoteSpec {
1355 token: &outbound_token,
1356 id: None,
1357 kind: "observation",
1358 name: Some("shared subject"),
1359 content,
1360 properties: Some(serde_json::json!({"direction": "outbound"})),
1361 },
1362 AtomicNoteSpec {
1363 token: &inbound_token,
1364 id: None,
1365 kind: "observation",
1366 name: Some("shared subject"),
1367 content,
1368 properties: Some(serde_json::json!({"direction": "inbound"})),
1369 },
1370 ],
1371 )
1372 .await
1373 })
1374 .await
1375 .expect("atomic note pair");
1376
1377 assert_eq!(notes.len(), 2);
1378 assert_eq!(notes[0].content, content);
1379 assert_eq!(notes[1].content, content);
1380 assert_eq!(
1381 usage.snapshot()["embed_calls"],
1382 2,
1383 "two identical notes across two models must issue one embed per model"
1384 );
1385 for model in ["atomic-dedup-model-a", "atomic-dedup-model-b"] {
1386 for (direction, token) in [("outbound", &outbound_token), ("inbound", &inbound_token)] {
1387 assert_eq!(
1388 runtime
1389 .vectors_for_model(token, model)
1390 .expect("vector store")
1391 .count()
1392 .await
1393 .expect("vector count"),
1394 1,
1395 "the {direction} note must retain its vector row for {model}"
1396 );
1397 }
1398 }
1399 }
1400
1401 #[tokio::test]
1406 async fn version_create_notes_atomic_upserts_fts_document_for_reused_note_id() {
1407 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
1408 let ns = "atomic-message-fts-upsert-test";
1409 let token = runtime
1410 .authorize(Namespace::parse(ns).unwrap())
1411 .expect("authorize");
1412 let id = Uuid::new_v4();
1413
1414 let first = create_notes_atomic(
1415 &runtime,
1416 vec![AtomicNoteSpec {
1417 token: &token,
1418 id: Some(id),
1419 kind: "observation",
1420 name: None,
1421 content: "first content",
1422 properties: None,
1423 }],
1424 )
1425 .await
1426 .expect("first write with supplied id");
1427 assert_eq!(first[0].version, 1);
1428
1429 let second = create_notes_atomic(
1430 &runtime,
1431 vec![AtomicNoteSpec {
1432 token: &token,
1433 id: Some(id),
1434 kind: "observation",
1435 name: None,
1436 content: "second content replacing the first",
1437 properties: None,
1438 }],
1439 )
1440 .await
1441 .expect("second write reusing the same id");
1442 assert_eq!(second[0].version, 2);
1443 assert_eq!(
1444 runtime
1445 .notes(&token)
1446 .unwrap()
1447 .get_note(id)
1448 .await
1449 .unwrap()
1450 .unwrap()
1451 .version,
1452 2
1453 );
1454 let repeated = create_notes_atomic(
1455 &runtime,
1456 (0..2)
1457 .map(|_| AtomicNoteSpec {
1458 token: &token,
1459 id: Some(id),
1460 kind: "observation",
1461 name: None,
1462 content: "second content replacing the first",
1463 properties: None,
1464 })
1465 .collect(),
1466 )
1467 .await
1468 .expect("equal-value upserts in the same batch");
1469 assert_eq!(
1470 repeated.iter().map(|note| note.version).collect::<Vec<_>>(),
1471 vec![3, 4]
1472 );
1473 assert_eq!(
1474 runtime
1475 .notes(&token)
1476 .unwrap()
1477 .get_note(id)
1478 .await
1479 .unwrap()
1480 .unwrap()
1481 .version,
1482 4
1483 );
1484
1485 assert_eq!(
1486 fts_row_count(&runtime, ns).await,
1487 1,
1488 "exactly one FTS document may exist for the reused id"
1489 );
1490
1491 let text = runtime.text_for_notes(&token).expect("text store");
1492 let doc = text
1493 .get_document(ns, id)
1494 .await
1495 .expect("get_document")
1496 .expect("document exists for the reused id");
1497 assert!(
1498 doc.body.contains("second content"),
1499 "the surviving FTS document must reflect the latest content; got {:?}",
1500 doc.body
1501 );
1502 }
1503
1504 #[tokio::test]
1518 async fn create_notes_atomic_derives_from_actor_overwriting_a_forged_value() {
1519 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
1520 runtime.install_note_write_validator(std::sync::Arc::new(stamp_from_actor));
1521 let ns = "atomic-message-identity-guard-test";
1522 let token = runtime
1523 .authorize(Namespace::parse(ns).unwrap())
1524 .expect("authorize");
1525 let true_actor = token.actor().id.clone();
1526
1527 let notes = create_notes_atomic(
1528 &runtime,
1529 vec![AtomicNoteSpec {
1530 token: &token,
1531 id: None,
1532 kind: "observation",
1533 name: None,
1534 content: "atomic writer identity guard probe",
1535 properties: Some(serde_json::json!({"from_actor": "forged-actor"})),
1536 }],
1537 )
1538 .await
1539 .expect("atomic note write must succeed — the guard derives, it does not refuse");
1540
1541 assert_eq!(notes.len(), 1);
1542 assert_eq!(
1543 notes[0]
1544 .properties
1545 .as_ref()
1546 .and_then(|p| p.get("from_actor"))
1547 .and_then(|v| v.as_str()),
1548 Some(true_actor.as_str()),
1549 "the atomic writer must store the token-derived from_actor, not the \
1550 caller-supplied forged value"
1551 );
1552 }
1553}