1#[cfg(test)]
2use super::race_seam;
3use super::{
4 entity_fts_document, map_merge_entity_storage_error, merge_entity_sql, merge_properties,
5 stale_entity_snapshot_error, Any, AtomicUnitOp, ContentMergeStrategy, EmbeddingModelPlan,
6 Entity, EntityDedupMergePolicy, EntityMergeValidation, EntityPatch, EventAttribution,
7 EventKind, HashMap, KhiveRuntime, MergeEventContext, MergeSqlError, MergeSummary,
8 MergeTxLimits, NamespaceToken, RuntimeError, RuntimeResult, RuntimeWriteOperation,
9 SqlStatement, SqlValue, SqliteError, SubstrateKind, Uuid, Value,
10};
11
12use crate::secret_gate_finalizer::entity_admission::{
13 EntityAdmission, OrdinaryEntityUpdateContext,
14};
15use crate::EntityCandidateMutation;
16
17#[derive(Clone, Copy)]
18enum EntityUpdateRoute<'a> {
19 ReservationOnly,
20 OrdinaryConstructor(&'a OrdinaryEntityUpdateContext),
21}
22
23struct PreparedEntityUpdate {
24 entity: Entity,
25 reindex_required: bool,
26 changed_fields: Vec<&'static str>,
27 expected_updated_at: i64,
28 expected_deleted_at: Option<i64>,
29 original: Entity,
30}
31
32impl KhiveRuntime {
33 pub(crate) async fn prepare_update_entity(
45 &self,
46 token: &NamespaceToken,
47 id: Uuid,
48 patch: EntityPatch,
49 ) -> RuntimeResult<(Entity, bool, Vec<&'static str>, i64, Option<i64>)> {
50 self.prepare_guarded_entity_update(token, id, patch, None, &[])
51 .await
52 }
53
54 async fn prepare_guarded_entity_update(
55 &self,
56 token: &NamespaceToken,
57 id: Uuid,
58 patch: EntityPatch,
59 expected: Option<&Entity>,
60 remove_properties: &[&str],
61 ) -> RuntimeResult<(Entity, bool, Vec<&'static str>, i64, Option<i64>)> {
62 let prepared = self
63 .prepare_entity_update_for_route(
64 token,
65 id,
66 patch,
67 expected,
68 remove_properties,
69 EntityUpdateRoute::ReservationOnly,
70 )
71 .await?;
72 Ok((
73 prepared.entity,
74 prepared.reindex_required,
75 prepared.changed_fields,
76 prepared.expected_updated_at,
77 prepared.expected_deleted_at,
78 ))
79 }
80
81 async fn prepare_entity_update_for_route(
82 &self,
83 token: &NamespaceToken,
84 id: Uuid,
85 patch: EntityPatch,
86 expected: Option<&Entity>,
87 remove_properties: &[&str],
88 route: EntityUpdateRoute<'_>,
89 ) -> RuntimeResult<PreparedEntityUpdate> {
90 crate::secret_gate::reject_reserved_secret_gate_property(patch.properties.as_ref())?;
91 if !remove_properties.is_empty() {
92 let removals = Value::Object(
93 remove_properties
94 .iter()
95 .map(|key| ((*key).to_string(), Value::Null))
96 .collect(),
97 );
98 crate::secret_gate::reject_reserved_secret_gate_property(Some(&removals))?;
99 }
100 if matches!(route, EntityUpdateRoute::ReservationOnly) {
101 if let Some(ref name) = patch.name {
102 crate::secret_gate::check_at(name, "entity", "name")?;
103 }
104 if let Some(Some(ref desc)) = patch.description {
105 crate::secret_gate::check_at(desc, "entity", "description")?;
106 }
107 if let Some(ref props) = patch.properties {
108 crate::secret_gate::check_json_at(props, "entity", "properties")?;
109 }
110 if let Some(ref tags) = patch.tags {
111 crate::secret_gate::check_tags_at(tags, "entity", "tags")?;
112 }
113 } else if let EntityUpdateRoute::OrdinaryConstructor(context) = route {
114 context.check_patch(
115 patch.name.as_deref(),
116 patch
117 .description
118 .as_ref()
119 .and_then(|value| value.as_deref()),
120 patch.properties.as_ref(),
121 patch.tags.as_deref(),
122 )?;
123 }
124 let store = self.entities(token)?;
125 let mut entity = store.get_entity(id).await?.ok_or_else(|| {
126 if expected.is_some() {
127 stale_entity_snapshot_error(id)
128 } else {
129 RuntimeError::NotFound(format!("entity {id}"))
130 }
131 })?;
132 if let Some(expected) = expected {
133 let actual = serde_json::to_value(&entity)
134 .map_err(|error| RuntimeError::Internal(error.to_string()))?;
135 let expected = serde_json::to_value(expected)
136 .map_err(|error| RuntimeError::Internal(error.to_string()))?;
137 if actual != expected {
138 return Err(stale_entity_snapshot_error(id));
139 }
140 }
141 let original = entity.clone();
142 let expected_updated_at = entity.updated_at;
143 let expected_deleted_at = entity.deleted_at;
144 #[cfg(test)]
145 race_seam::pause_after_read().await;
146
147 let validated_entity_type = match &patch.entity_type {
151 Some(None) => Some(None),
152 Some(Some(raw)) => Some(Some(
153 self.validate_entity_type_for_kind(&entity.kind, Some(raw))?
154 .expect("set branch always yields a normalized value"),
155 )),
156 None => None,
157 };
158
159 let mut reindex_required = false;
160 let mut changed_fields: Vec<&'static str> = Vec::new();
161
162 if let Some(name) = patch.name {
163 reindex_required |= entity.name != name;
164 entity.name = name;
165 changed_fields.push("name");
166 }
167 if let Some(desc_patch) = patch.description {
168 reindex_required |= entity.description != desc_patch;
169 entity.description = desc_patch;
170 changed_fields.push("description");
171 }
172 if let Some(props) = patch.properties {
173 let (merged, _) = merge_properties(
174 &entity.properties,
175 &Some(props),
176 EntityDedupMergePolicy::PreferFrom,
177 );
178 entity.properties = merged;
179 changed_fields.push("properties");
180 }
181 if let Some(Value::Object(properties)) = entity.properties.as_mut() {
182 let mut removed = false;
183 for key in remove_properties {
184 removed |= properties.remove(*key).is_some();
185 }
186 if removed && !changed_fields.contains(&"properties") {
187 changed_fields.push("properties");
188 }
189 }
190 if let Some(tags) = patch.tags {
191 entity.tags = tags;
192 changed_fields.push("tags");
193 }
194 if let Some(entity_type) = validated_entity_type {
195 reindex_required |= entity.entity_type != entity_type;
196 entity.entity_type = entity_type;
197 changed_fields.push("entity_type");
198 }
199
200 crate::secret_gate::reject_reserved_secret_gate_property(entity.properties.as_ref())?;
203
204 if expected.is_some() && changed_fields.is_empty() {
205 return Ok(PreparedEntityUpdate {
206 entity,
207 reindex_required,
208 changed_fields,
209 expected_updated_at,
210 expected_deleted_at,
211 original,
212 });
213 }
214
215 if let Some(hook) = self.entity_kind_hook(&entity.kind) {
224 hook.validate_entity_update(self, token, &entity, entity.properties.as_ref())
225 .await?;
226 }
227
228 let minimum_updated_at = expected_updated_at.checked_add(1).ok_or_else(|| {
234 RuntimeError::Internal(format!(
235 "entity {id} updated_at is already at i64::MAX and cannot advance"
236 ))
237 })?;
238 entity.updated_at = chrono::Utc::now()
239 .timestamp_micros()
240 .max(minimum_updated_at);
241 Ok(PreparedEntityUpdate {
242 entity,
243 reindex_required,
244 changed_fields,
245 expected_updated_at,
246 expected_deleted_at,
247 original,
248 })
249 }
250
251 #[cfg(test)]
252 pub(crate) async fn update_entity(
253 &self,
254 token: &NamespaceToken,
255 id: Uuid,
256 patch: EntityPatch,
257 ) -> RuntimeResult<Entity> {
258 Ok(self
259 .update_entity_with_embedding_report(token, id, patch)
260 .await?
261 .0)
262 }
263
264 pub async fn update_entity_with_embedding_report(
265 &self,
266 token: &NamespaceToken,
267 id: Uuid,
268 patch: EntityPatch,
269 ) -> RuntimeResult<(Entity, crate::retrieval::EmbeddingTruncationReport)> {
270 self.update_entity_with_expected_version_and_embedding_report(token, id, patch, None)
271 .await
272 }
273
274 pub async fn update_entity_with_expected_version_and_embedding_report(
276 &self,
277 token: &NamespaceToken,
278 id: Uuid,
279 patch: EntityPatch,
280 expected_version: Option<i64>,
281 ) -> RuntimeResult<(Entity, crate::retrieval::EmbeddingTruncationReport)> {
282 crate::entity_write::validate_expected_version(expected_version)?;
283 let admission = OrdinaryEntityUpdateContext::capture(token);
284 let PreparedEntityUpdate {
285 entity,
286 reindex_required,
287 changed_fields,
288 expected_updated_at,
289 expected_deleted_at,
290 original,
291 } = self
292 .prepare_entity_update_for_route(
293 token,
294 id,
295 patch,
296 None,
297 &[],
298 EntityUpdateRoute::OrdinaryConstructor(&admission),
299 )
300 .await?;
301 let entity = match admission.admit(token, entity)? {
302 EntityAdmission::Legacy(entity) => entity,
303 EntityAdmission::Exempt(prepared) => {
304 use crate::atomic_plan::PostCommitEffect;
305 use crate::atomic_runner::{
306 run_atomic_unit, AtomicOpFailure, AtomicOpPlan, AtomicRunOutcome,
307 };
308 let mut entity = prepared.entity().clone();
309 let next_version = original
310 .version
311 .checked_add(1)
312 .ok_or_else(|| RuntimeError::InvalidInput("entity version overflow".into()))?;
313 let (mut required, report) = self
314 .prepare_admitted_entity_indexes(token, &entity, &[], reindex_required)
315 .await?;
316 required.extend(crate::atomic_prepare::event_append_statements(
317 token, &entity.namespace, "update", EventKind::EntityUpdated,
318 SubstrateKind::Entity, entity.id,
319 serde_json::json!({"id": entity.id, "namespace": entity.namespace, "changed_fields": changed_fields}),
320 )?);
321 let plan = prepared.into_plan(
322 EntityCandidateMutation::ReplaceIfUnchanged {
323 expected: original,
324 expected_version,
325 },
326 required,
327 PostCommitEffect::None,
328 )?;
329 match run_atomic_unit(
330 self.sql().as_ref(),
331 vec![AtomicOpPlan::FinalizeEntity(Box::new(plan))],
332 )
333 .await
334 {
335 Ok(AtomicRunOutcome::Committed { .. }) => {
336 entity.version = next_version;
337 return Ok((entity, report));
338 }
339 Ok(AtomicRunOutcome::RolledBack {
340 failure: AtomicOpFailure::EntityConflict(conflict),
341 ..
342 }) => {
343 return Err(conflict.into_error().into());
344 }
345 Ok(AtomicRunOutcome::RolledBack {
346 failure: AtomicOpFailure::GuardFailed { .. },
347 ..
348 }) => {
349 return Err(stale_entity_snapshot_error(id));
350 }
351 Ok(AtomicRunOutcome::RolledBack { failure, .. }) => {
352 return Err(RuntimeError::Internal(format!(
353 "entity update finalization rolled back: {failure:?}"
354 )));
355 }
356 Err(error) => return Err(RuntimeError::Storage(error.0)),
357 }
358 }
359 };
360
361 self.persist_prepared_entity_update(
362 token,
363 entity,
364 reindex_required,
365 changed_fields,
366 expected_updated_at,
367 expected_deleted_at,
368 expected_version,
369 )
370 .await
371 }
372
373 pub async fn update_entity_if_unchanged(
380 &self,
381 token: &NamespaceToken,
382 expected: &Entity,
383 patch: EntityPatch,
384 remove_properties: &[&str],
385 ) -> RuntimeResult<Entity> {
386 let (entity, embedding) = self
387 .update_entity_if_unchanged_with_embedding_report(
388 token,
389 expected,
390 patch,
391 remove_properties,
392 )
393 .await?;
394 crate::operations::legacy_post_commit_result_with_embedding(
395 "update_entity_if_unchanged",
396 entity.id,
397 entity,
398 embedding,
399 Vec::new(),
400 )
401 }
402
403 pub async fn update_entity_if_unchanged_with_embedding_report(
405 &self,
406 token: &NamespaceToken,
407 expected: &Entity,
408 patch: EntityPatch,
409 remove_properties: &[&str],
410 ) -> RuntimeResult<(Entity, crate::retrieval::EmbeddingTruncationReport)> {
411 let (entity, reindex_required, changed_fields, expected_updated_at, expected_deleted_at) =
412 self.prepare_guarded_entity_update(
413 token,
414 expected.id,
415 patch,
416 Some(expected),
417 remove_properties,
418 )
419 .await?;
420 if changed_fields.is_empty() {
421 return Ok((
422 entity,
423 crate::retrieval::EmbeddingTruncationReport::default(),
424 ));
425 }
426 self.persist_prepared_entity_update(
427 token,
428 entity,
429 reindex_required,
430 changed_fields,
431 expected_updated_at,
432 expected_deleted_at,
433 None,
434 )
435 .await
436 }
437
438 #[allow(clippy::too_many_arguments)]
439 pub(super) async fn persist_prepared_entity_update(
440 &self,
441 token: &NamespaceToken,
442 mut entity: Entity,
443 reindex_required: bool,
444 changed_fields: Vec<&'static str>,
445 expected_updated_at: i64,
446 expected_deleted_at: Option<i64>,
447 expected_version: Option<i64>,
448 ) -> RuntimeResult<(Entity, crate::retrieval::EmbeddingTruncationReport)> {
449 crate::secret_gate::reject_reserved_secret_gate_property(entity.properties.as_ref())?;
452 let id = entity.id;
453 let _ = self.entities(token)?;
454 let next_version = entity
455 .version
456 .checked_add(1)
457 .ok_or_else(|| RuntimeError::InvalidInput("entity version overflow".into()))?;
458 use crate::atomic_plan::{AffectedRowGuard, PlanStatement, PostCommitEffect, UpdatePlan};
459 use crate::atomic_runner::{
460 run_atomic_unit, AtomicOpFailure, AtomicOpPlan, AtomicRunOutcome,
461 };
462 let plan = UpdatePlan {
463 target_id: id,
464 statements: vec![PlanStatement {
465 statement: khive_db::stores::entity::entity_replace_if_unchanged_statement(
466 &entity,
467 expected_updated_at,
468 expected_deleted_at,
469 ),
470 guard: Some(AffectedRowGuard::exactly(1)),
471 }],
472 post_commit: PostCommitEffect::None,
473 edge_natural_key: None,
474 idempotent_noop: false,
475 entity_guard: expected_version.map(|expected_version| {
476 crate::entity_write::EntityWriteGuard {
477 id,
478 expected_version,
479 }
480 }),
481 note_guard: None,
482 note_vector_purge: None,
483 note_embedding_inheritance: None,
484 graph_effects: Vec::new(),
485 };
486 match run_atomic_unit(
487 self.sql().as_ref(),
488 vec![AtomicOpPlan::Update(Box::new(plan))],
489 )
490 .await
491 {
492 Ok(AtomicRunOutcome::Committed { .. }) => entity.version = next_version,
493 Ok(AtomicRunOutcome::RolledBack {
494 failure: AtomicOpFailure::EntityConflict(conflict),
495 ..
496 }) => return Err(conflict.into_error().into()),
497 Ok(AtomicRunOutcome::RolledBack {
498 failure: AtomicOpFailure::GuardFailed { .. },
499 ..
500 }) => return Err(stale_entity_snapshot_error(id)),
501 Ok(AtomicRunOutcome::RolledBack { failure, .. }) => {
502 return Err(RuntimeError::Internal(format!(
503 "entity update rolled back: {failure:?}"
504 )))
505 }
506 Err(error) => return Err(RuntimeError::Storage(error.0)),
507 }
508
509 let event_token =
510 token.with_namespace(crate::Namespace::parse(&entity.namespace).map_err(|error| {
511 RuntimeError::Internal(format!("entity namespace invalid: {error}"))
512 })?);
513 let event_store = self.events(&event_token)?;
514 let event = khive_storage::event::Event::new(
515 entity.namespace.clone(),
516 "update",
517 EventKind::EntityUpdated,
518 SubstrateKind::Entity,
519 "",
520 )
521 .with_target(entity.id)
522 .with_payload(serde_json::json!({
523 "id": entity.id,
524 "namespace": entity.namespace,
525 "changed_fields": changed_fields,
526 }));
527 let event_result = event_store.append_event(event).await.map_err(|e| {
528 RuntimeError::Internal(format!("update_entity: event store write failed: {e}"))
529 });
530
531 let embedding_report = if reindex_required {
532 self.reindex_entity(token, &entity).await?
533 } else {
534 crate::retrieval::EmbeddingTruncationReport::default()
535 };
536
537 event_result?;
538
539 Ok((entity, embedding_report))
540 }
541
542 pub async fn merge_entity(
556 &self,
557 token: &NamespaceToken,
558 into_id: Uuid,
559 from_id: Uuid,
560 strategy: EntityDedupMergePolicy,
561 content_strategy: ContentMergeStrategy,
562 dry_run: bool,
563 ) -> RuntimeResult<MergeSummary> {
564 self.merge_entity_with_reason(
565 token,
566 into_id,
567 from_id,
568 strategy,
569 content_strategy,
570 dry_run,
571 None,
572 )
573 .await
574 }
575
576 #[allow(clippy::too_many_arguments)]
580 pub async fn merge_entity_with_reason(
581 &self,
582 token: &NamespaceToken,
583 into_id: Uuid,
584 from_id: Uuid,
585 strategy: EntityDedupMergePolicy,
586 content_strategy: ContentMergeStrategy,
587 dry_run: bool,
588 reason: Option<String>,
589 ) -> RuntimeResult<MergeSummary> {
590 self.merge_entity_with_validation(
591 token,
592 into_id,
593 from_id,
594 strategy,
595 content_strategy,
596 dry_run,
597 reason,
598 EntityMergeValidation::LegacyKind,
599 )
600 .await
601 }
602
603 #[allow(clippy::too_many_arguments)]
612 pub async fn merge_entity_with_reason_and_force(
613 &self,
614 token: &NamespaceToken,
615 into_id: Uuid,
616 from_id: Uuid,
617 strategy: EntityDedupMergePolicy,
618 content_strategy: ContentMergeStrategy,
619 dry_run: bool,
620 reason: Option<String>,
621 force: bool,
622 ) -> RuntimeResult<MergeSummary> {
623 let validation = if force {
624 EntityMergeValidation::Forced
625 } else {
626 EntityMergeValidation::SafetyFloor
627 };
628 self.merge_entity_with_validation(
629 token,
630 into_id,
631 from_id,
632 strategy,
633 content_strategy,
634 dry_run,
635 reason,
636 validation,
637 )
638 .await
639 }
640
641 #[allow(clippy::too_many_arguments)]
642 async fn merge_entity_with_validation(
643 &self,
644 token: &NamespaceToken,
645 into_id: Uuid,
646 from_id: Uuid,
647 strategy: EntityDedupMergePolicy,
648 content_strategy: ContentMergeStrategy,
649 dry_run: bool,
650 reason: Option<String>,
651 validation: EntityMergeValidation,
652 ) -> RuntimeResult<MergeSummary> {
653 if let Some(reason) = reason.as_deref() {
654 crate::secret_gate::check_at(reason, "merge", "reason")?;
655 }
656 if into_id == from_id {
657 return Err(RuntimeError::InvalidInput(
658 "cannot merge an entity into itself".into(),
659 ));
660 }
661 let ns = token.namespace().as_str().to_owned();
662 let fts_table = "fts_entities".to_string();
663 let embedding_plan = EmbeddingModelPlan::capture(self);
667 let vec_tables = embedding_plan.vector_tables();
668 let pack_rules = self.pack_edge_rules();
671
672 let _ = self.entities(token)?;
674 let _ = self.graph(token)?;
675 let _ = self.text(token)?;
676 let _ = self.events(token)?;
677 for model_name in embedding_plan.model_names() {
680 let _ = self.vectors_for_model(token, model_name)?;
681 }
682
683 let pool = self.backend().pool_arc();
684 let writer_task = pool
685 .writer_task_for_runtime_write(RuntimeWriteOperation::MergeEntity)
686 .map_err(RuntimeError::Storage)?;
687 let merge_event_id = Uuid::new_v4();
690 let event_context = MergeEventContext {
691 attribution: EventAttribution::from_token(token),
692 reason,
693 force: validation == EntityMergeValidation::Forced,
694 strategy,
695 content_strategy,
696 kind: EventKind::EntityMerged,
697 substrate: SubstrateKind::Entity,
698 event_id: Some(merge_event_id),
699 };
700
701 let (mut summary, updated_entity) = if let Some(writer_task) = writer_task {
702 writer_task
703 .send(move |conn| {
704 merge_entity_sql(
705 conn,
706 ns,
707 fts_table,
708 vec_tables,
709 into_id,
710 from_id,
711 strategy,
712 content_strategy,
713 dry_run,
714 pack_rules,
715 validation,
716 MergeTxLimits::default(),
717 merge_event_id,
718 Some(event_context),
719 )
720 .map_err(|e| {
721 e.into_storage_error(
722 khive_storage::StorageCapability::Entities,
723 "merge_entity",
724 )
725 })
726 })
727 .await
728 .inspect_err(|error| khive_storage::usage::account_event_write(Err(error)))
729 .map_err(map_merge_entity_storage_error)?
730 } else {
731 tokio::task::spawn_blocking(move || {
732 let guard = pool.writer()?;
733 let mut refusal = None;
734 let result = guard.transaction(|conn| {
735 merge_entity_sql(
736 conn,
737 ns,
738 fts_table,
739 vec_tables,
740 into_id,
741 from_id,
742 strategy,
743 content_strategy,
744 dry_run,
745 pack_rules,
746 validation,
747 MergeTxLimits::default(),
748 merge_event_id,
749 Some(event_context),
750 )
751 .map_err(|error| match error {
752 MergeSqlError::Sqlite(error) => error,
753 MergeSqlError::Refusal(error) => {
754 refusal = Some(error);
755 SqliteError::InvalidData(
756 "entity merge refused by transactional policy".to_string(),
757 )
758 }
759 })
760 });
761 match refusal {
762 Some(error) => Err(error),
763 None => result.map_err(RuntimeError::from),
764 }
765 })
766 .await
767 .map_err(|e| RuntimeError::Internal(e.to_string()))??
768 };
769
770 if !dry_run {
772 khive_storage::usage::account_event_write(Ok(1));
773 tracing::info!(
774 into_id = %summary.kept_id,
775 from_id = %summary.removed_id,
776 budget_rows = summary.tx_budget.rows_charged,
777 budget_bytes = summary.tx_budget.bytes_charged,
778 budget_max_rows = summary.tx_budget.max_rows,
779 budget_max_bytes = summary.tx_budget.max_bytes,
780 "merge_entity: transaction materialization budget"
781 );
782 }
783
784 if !dry_run && !embedding_plan.is_empty() {
787 match self
788 .reindex_entity_with_plan(token, &updated_entity, &embedding_plan, None)
789 .await
790 {
791 Ok(report) => summary.embedding_truncation = report,
792 Err(error) => {
793 tracing::warn!(
794 into_id = %summary.kept_id,
795 from_id = %summary.removed_id,
796 error = %error,
797 "merge_entity: committed merge but survivor reindex failed"
798 );
799 summary.post_commit_reindex_error = Some(error.to_string());
800 }
801 }
802 }
803
804 Ok(summary)
805 }
806
807 async fn apply_entity_index_revision(
810 &self,
811 entity: &Entity,
812 statements: Vec<SqlStatement>,
813 ) -> RuntimeResult<bool> {
814 let namespace = entity.namespace.clone();
815 let id = entity.id.to_string();
816 let version = entity.version;
817 let op: AtomicUnitOp = Box::new(move |writer| {
818 Box::pin(async move {
819 let current = writer
820 .query_scalar(SqlStatement {
821 sql: "SELECT version FROM entities \
822 WHERE namespace=?1 AND id=?2 AND deleted_at IS NULL"
823 .into(),
824 params: vec![SqlValue::Text(namespace), SqlValue::Text(id)],
825 label: Some("entity-index-revision".into()),
826 })
827 .await?;
828 if !matches!(current, Some(SqlValue::Integer(current)) if current == version) {
829 return Ok(Box::new(false) as Box<dyn Any + Send>);
830 }
831 for statement in statements {
832 writer.execute(statement).await?;
833 }
834 Ok(Box::new(true) as Box<dyn Any + Send>)
835 })
836 });
837 self.sql()
838 .atomic_unit(op)
839 .await?
840 .downcast::<bool>()
841 .map(|result| *result)
842 .map_err(|_| RuntimeError::Internal("invalid entity index outcome".into()))
843 }
844
845 pub(crate) fn entity_vector_insert_statements(
846 table: &str,
847 entity: &Entity,
848 model_name: &str,
849 vector: &[f32],
850 ) -> Vec<SqlStatement> {
851 let subject = entity.id.to_string();
852 let model_key = table
853 .strip_prefix("vec_")
854 .expect("runtime vector tables use the vec_ prefix");
855 let kind = SubstrateKind::Entity.to_string();
856 let field = "entity.body";
857 let blob = khive_storage::encode_f32_native(vector);
858 vec![
859 SqlStatement {
860 sql: format!(
861 "INSERT INTO ann_write_log \
862 (namespace, embedding_model, kind, field, subject_id, op) \
863 SELECT namespace, embedding_model, kind, field, subject_id, 'delete' \
864 FROM {table} WHERE subject_id=?1 AND NOT \
865 (namespace=?2 AND embedding_model=?3 AND kind=?4 AND field=?5)"
866 ),
867 params: vec![
868 SqlValue::Text(subject.clone()),
869 SqlValue::Text(entity.namespace.clone()),
870 SqlValue::Text(model_name.to_string()),
871 SqlValue::Text(kind.clone()),
872 SqlValue::Text(field.into()),
873 ],
874 label: Some("entity-reindex-log-delete".into()),
875 },
876 SqlStatement {
877 sql: format!("DELETE FROM {table} WHERE subject_id=?1"),
878 params: vec![SqlValue::Text(subject.clone())],
879 label: Some("entity-reindex-vector-delete".into()),
880 },
881 SqlStatement {
885 sql: "DELETE FROM vector_provenance \
886 WHERE model_key = ?1 AND subject_id = ?2"
887 .into(),
888 params: vec![
889 SqlValue::Text(model_key.to_string()),
890 SqlValue::Text(subject.clone()),
891 ],
892 label: Some("entity-reindex-provenance-clear".into()),
893 },
894 SqlStatement {
895 sql: format!(
896 "INSERT INTO {table} \
897 (subject_id, namespace, kind, field, embedding_model, embedding) \
898 VALUES (?1, ?2, ?3, ?4, ?5, ?6)"
899 ),
900 params: vec![
901 SqlValue::Text(subject.clone()),
902 SqlValue::Text(entity.namespace.clone()),
903 SqlValue::Text(kind.clone()),
904 SqlValue::Text(field.into()),
905 SqlValue::Text(model_name.to_string()),
906 SqlValue::Blob(blob),
907 ],
908 label: Some("entity-reindex-vector-insert".into()),
909 },
910 SqlStatement {
911 sql: "INSERT INTO ann_write_log \
912 (namespace, embedding_model, kind, field, subject_id, op) \
913 VALUES (?1, ?2, ?3, ?4, ?5, 'upsert')"
914 .into(),
915 params: vec![
916 SqlValue::Text(entity.namespace.clone()),
917 SqlValue::Text(model_name.to_string()),
918 SqlValue::Text(kind),
919 SqlValue::Text(field.into()),
920 SqlValue::Text(subject),
921 ],
922 label: Some("entity-reindex-log-upsert".into()),
923 },
924 ]
925 }
926
927 pub(crate) async fn publish_entity_vector_revision(
928 &self,
929 token: &NamespaceToken,
930 entity: &Entity,
931 model_name: &str,
932 vector: &[f32],
933 ) -> RuntimeResult<bool> {
934 self.vectors_for_model(token, model_name)?;
935 let (storage_model, dimensions) = self.vector_model_metadata(model_name)?;
936 if let Some(index) = crate::atomic_message::non_finite_index(vector) {
937 return Err(RuntimeError::InvalidInput(format!(
938 "non-finite entity vector at index {index}"
939 )));
940 }
941 if vector.len() != dimensions {
942 return Err(RuntimeError::InvalidInput(format!(
943 "entity vector has {} dimensions; expected {dimensions}",
944 vector.len()
945 )));
946 }
947 let table = format!("vec_{}", crate::config::sanitize_key(&storage_model));
948 let statements =
949 Self::entity_vector_insert_statements(&table, entity, &storage_model, vector);
950 #[cfg(test)]
951 race_seam::pause_before_entity_vector_publish().await;
952 self.apply_entity_index_revision(entity, statements).await
953 }
954
955 pub(crate) async fn reindex_entity(
968 &self,
969 token: &NamespaceToken,
970 entity: &Entity,
971 ) -> RuntimeResult<crate::retrieval::EmbeddingTruncationReport> {
972 let embedding_plan = EmbeddingModelPlan::capture(self);
973 self.reindex_entity_with_plan(token, entity, &embedding_plan, None)
974 .await
975 }
976
977 pub(crate) async fn reindex_entity_with_precomputed(
978 &self,
979 token: &NamespaceToken,
980 entity: &Entity,
981 mut precomputed: HashMap<String, crate::retrieval::DocumentEmbeddingOutcome>,
982 ) -> RuntimeResult<crate::retrieval::EmbeddingTruncationReport> {
983 let embedding_plan = EmbeddingModelPlan::capture(self);
984 self.reindex_entity_with_plan(token, entity, &embedding_plan, Some(&mut precomputed))
985 .await
986 }
987
988 pub(super) async fn reindex_entity_with_plan(
989 &self,
990 token: &NamespaceToken,
991 entity: &Entity,
992 embedding_plan: &EmbeddingModelPlan,
993 mut precomputed: Option<&mut HashMap<String, crate::retrieval::DocumentEmbeddingOutcome>>,
994 ) -> RuntimeResult<crate::retrieval::EmbeddingTruncationReport> {
995 #[cfg(test)]
998 if crate::operations::consume_fts_fail_fault(&entity.namespace) {
999 return Err(RuntimeError::Internal("injected FTS failure".to_string()));
1000 }
1001 let doc = entity_fts_document(entity);
1003 let embed_body = doc.body.clone();
1004 let _ = self.text(token)?;
1005 #[cfg(test)]
1006 race_seam::pause_before_entity_index_publish().await;
1007 let statements = khive_db::stores::text::delete_document_statements(
1008 "fts_entities",
1009 &entity.namespace,
1010 entity.id,
1011 )
1012 .into_iter()
1013 .chain(khive_db::stores::text::insert_document_statements(
1014 "fts_entities",
1015 &doc,
1016 ))
1017 .collect();
1018 if !self.apply_entity_index_revision(entity, statements).await? {
1019 return Ok(crate::retrieval::EmbeddingTruncationReport::default());
1020 }
1021
1022 let mut report = crate::retrieval::EmbeddingTruncationReport::default();
1023 for model_name in embedding_plan.model_names() {
1024 let embedding = match precomputed
1025 .as_mut()
1026 .and_then(|outcomes| outcomes.remove(model_name))
1027 {
1028 Some(outcome) => Ok(outcome),
1029 None => {
1030 self.embed_document_with_model_outcome_for_token(token, model_name, &embed_body)
1031 .await
1032 }
1033 };
1034 match embedding {
1035 Ok(outcome) => {
1036 report.observe(&outcome);
1037 match self
1038 .publish_entity_vector_revision(token, entity, model_name, &outcome.vector)
1039 .await
1040 {
1041 Ok(true) => {}
1042 Ok(false) => break,
1043 Err(error) => tracing::warn!(
1044 model = model_name,
1045 id = %entity.id,
1046 "reindex_entity: vector insert failed, skipping model: {error}"
1047 ),
1048 }
1049 }
1050 Err(e) => {
1051 tracing::warn!(
1052 model = model_name,
1053 id = %entity.id,
1054 "reindex_entity: embed failed for model, skipping: {e}"
1055 );
1056 }
1057 }
1058 }
1059
1060 Ok(report)
1061 }
1062}