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