1use std::collections::{HashMap, HashSet};
7
8use anyhow::Context;
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11use uuid::Uuid;
12
13use khive_storage::types::{EdgeFilter, LinkId, PageRequest};
14use khive_storage::{EdgeRelation, EntityFilter};
15use khive_types::{EntityKind, ImportKindPolicy};
16use std::str::FromStr;
17
18use crate::error::{RuntimeError, RuntimeResult};
19use crate::runtime::{KhiveRuntime, NamespaceToken};
20
21pub fn parse_archive_timestamp(
26 value: Option<&str>,
27 fallback: DateTime<Utc>,
28) -> anyhow::Result<DateTime<Utc>> {
29 let Some(raw) = value else {
30 return Ok(fallback);
31 };
32 DateTime::parse_from_rfc3339(raw)
33 .map(|dt| dt.with_timezone(&Utc))
34 .with_context(|| format!("timestamp {raw:?} must be RFC3339"))
35}
36
37#[derive(Clone, Debug, Serialize, Deserialize)]
44pub struct KgArchive {
45 pub format: String,
46 pub version: String,
47 pub namespace: String,
48 pub exported_at: DateTime<Utc>,
49 pub entities: Vec<ExportedEntity>,
50 pub edges: Vec<ExportedEdge>,
51}
52
53#[derive(Clone, Debug, Serialize, Deserialize)]
55pub struct ExportedEntity {
56 pub id: Uuid,
57 pub kind: String,
59 #[serde(skip_serializing_if = "Option::is_none")]
61 pub entity_type: Option<String>,
62 pub name: String,
63 #[serde(skip_serializing_if = "Option::is_none")]
64 pub description: Option<String>,
65 #[serde(skip_serializing_if = "Option::is_none")]
66 pub properties: Option<serde_json::Value>,
67 #[serde(default)]
68 pub tags: Vec<String>,
69 pub created_at: DateTime<Utc>,
70 pub updated_at: DateTime<Utc>,
71}
72
73#[derive(Clone, Debug, Serialize, Deserialize)]
75pub struct ExportedEdge {
76 #[serde(default = "Uuid::new_v4")]
81 pub edge_id: Uuid,
82 pub source: Uuid,
83 pub target: Uuid,
84 pub relation: EdgeRelation,
86 pub weight: f64,
87 #[serde(default, skip_serializing_if = "Option::is_none")]
89 pub properties: Option<serde_json::Value>,
90 #[serde(
93 default = "default_import_timestamp",
94 deserialize_with = "deserialize_import_timestamp"
95 )]
96 pub created_at: DateTime<Utc>,
97 #[serde(
100 default = "default_import_timestamp",
101 deserialize_with = "deserialize_import_timestamp"
102 )]
103 pub updated_at: DateTime<Utc>,
104}
105
106fn default_import_timestamp() -> DateTime<Utc> {
107 Utc::now()
108}
109
110fn deserialize_import_timestamp<'de, D>(deserializer: D) -> Result<DateTime<Utc>, D::Error>
111where
112 D: serde::Deserializer<'de>,
113{
114 let raw = String::deserialize(deserializer)?;
115 DateTime::parse_from_rfc3339(&raw)
116 .map(|timestamp| timestamp.with_timezone(&Utc))
117 .map_err(|error| serde::de::Error::custom(format!("timestamp must be RFC3339: {error}")))
118}
119
120#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
122#[serde(rename_all = "snake_case")]
123pub enum ImportEdgeSkipCause {
124 MissingSource,
126 MissingTarget,
128 EndpointContract,
130}
131
132#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
134pub struct ImportSkippedEdge {
135 pub edge_id: Uuid,
136 pub source: Uuid,
137 pub target: Uuid,
138 pub relation: EdgeRelation,
139 pub cause: ImportEdgeSkipCause,
140 #[serde(default, skip_serializing_if = "Option::is_none")]
142 pub error: Option<String>,
143}
144
145#[derive(Clone, Debug, Serialize, Deserialize)]
147pub struct ImportSummary {
148 pub entities_imported: usize,
149 pub edges_imported: usize,
150 pub edges_skipped: usize,
155 #[serde(default, skip_serializing_if = "Vec::is_empty")]
159 pub skipped_edges: Vec<ImportSkippedEdge>,
160 #[serde(default)]
163 pub embedding_truncation: crate::retrieval::EmbeddingTruncationReport,
164}
165
166#[derive(Clone, Debug, Default, PartialEq, Eq)]
169pub struct ImportValidation {
170 pub unknown_entity_kinds: Vec<String>,
171}
172
173impl KhiveRuntime {
176 pub async fn export_kg(&self, token: &NamespaceToken) -> RuntimeResult<KgArchive> {
182 let ns = token.namespace().as_str().to_owned();
183
184 let entity_page = self
185 .entities(token)?
186 .query_entities(
187 &ns,
188 EntityFilter::default(),
189 PageRequest {
190 offset: 0,
191 limit: u32::MAX,
192 },
193 )
194 .await?;
195
196 let entities: Vec<ExportedEntity> = entity_page
197 .items
198 .into_iter()
199 .map(|e| {
200 let created_at =
201 DateTime::from_timestamp_micros(e.created_at).unwrap_or_else(Utc::now);
202 let updated_at =
203 DateTime::from_timestamp_micros(e.updated_at).unwrap_or_else(Utc::now);
204 ExportedEntity {
205 id: e.id,
206 kind: e.kind.to_string(),
207 entity_type: e.entity_type,
208 name: e.name,
209 description: e.description,
210 properties: e.properties,
211 tags: e.tags,
212 created_at,
213 updated_at,
214 }
215 })
216 .collect();
217
218 let source_ids: Vec<Uuid> = entities.iter().map(|e| e.id).collect();
219 let edges = if source_ids.is_empty() {
220 Vec::new()
221 } else {
222 let filter = EdgeFilter {
223 source_ids: source_ids.clone(),
224 ..Default::default()
225 };
226 let edge_page = self
227 .graph(token)?
228 .query_edges(
229 filter,
230 Vec::new(),
231 PageRequest {
232 offset: 0,
233 limit: u32::MAX,
234 },
235 )
236 .await?;
237
238 let id_set: HashSet<Uuid> = source_ids.into_iter().collect();
239 edge_page
240 .items
241 .into_iter()
242 .filter(|e| id_set.contains(&e.source_id))
243 .map(|e| ExportedEdge {
244 edge_id: e.id.into(),
245 source: e.source_id,
246 target: e.target_id,
247 relation: e.relation,
248 weight: e.weight,
249 properties: e.metadata,
250 created_at: e.created_at,
251 updated_at: e.updated_at,
252 })
253 .collect()
254 };
255
256 Ok(KgArchive {
257 format: "khive-kg".to_string(),
258 version: "0.1".to_string(),
259 namespace: ns,
260 exported_at: Utc::now(),
261 entities,
262 edges,
263 })
264 }
265
266 pub async fn export_kg_json(&self, token: &NamespaceToken) -> RuntimeResult<String> {
268 let archive = self.export_kg(token).await?;
269 serde_json::to_string(&archive).map_err(|e| RuntimeError::InvalidInput(e.to_string()))
270 }
271
272 pub async fn import_kg(
281 &self,
282 archive: &KgArchive,
283 token: &NamespaceToken,
284 ) -> RuntimeResult<ImportSummary> {
285 self.import_kg_with_policy(archive, token, ImportKindPolicy::Strict)
286 .await
287 }
288
289 pub fn validate_kg_import(
291 &self,
292 archive: &KgArchive,
293 policy: ImportKindPolicy,
294 ) -> RuntimeResult<ImportValidation> {
295 if archive.format != "khive-kg" {
296 return Err(RuntimeError::InvalidInput(format!(
297 "unsupported archive format {:?}; expected \"khive-kg\"",
298 archive.format
299 )));
300 }
301 if archive.version != "0.1" {
302 return Err(RuntimeError::InvalidInput(format!(
303 "unsupported archive version {:?}; supported: \"0.1\"",
304 archive.version
305 )));
306 }
307
308 let kinds = self.import_entity_kind_registry()?;
312 let mut validation = ImportValidation::default();
313 let mut seen_unknown = HashSet::new();
314 for (index, entity) in archive.entities.iter().enumerate() {
315 if entity.kind.trim().is_empty() {
316 return Err(RuntimeError::InvalidInput(format!(
317 "archive entity {index} kind must be non-blank"
318 )));
319 }
320 let base_kind = EntityKind::from_str(&entity.kind).is_ok();
321 let normalized = entity.kind.trim().to_ascii_lowercase();
322 let known =
323 kinds.iter().any(|kind| kind == &entity.kind) || (kinds.is_empty() && base_kind);
324 if !known {
325 let recognized = base_kind
328 || kinds
329 .iter()
330 .any(|kind| kind.eq_ignore_ascii_case(&normalized));
331 if policy == ImportKindPolicy::Strict || recognized {
332 let valid = if kinds.is_empty() {
333 EntityKind::ALL
334 .iter()
335 .map(|kind| kind.name())
336 .collect::<Vec<_>>()
337 .join(", ")
338 } else {
339 kinds.join(", ")
340 };
341 return Err(RuntimeError::InvalidInput(format!(
342 "archive entity {index} ({}) has unknown entity kind {:?} (or non-canonical spelling); valid: {valid}",
343 entity.id, entity.kind
344 )));
345 }
346 if seen_unknown.insert(entity.kind.clone()) {
347 validation.unknown_entity_kinds.push(entity.kind.clone());
348 }
349 }
350 crate::secret_gate::reject_reserved_secret_gate_property(entity.properties.as_ref())?;
355 if entity.name.trim().is_empty() {
356 return Err(RuntimeError::InvalidInput(format!(
357 "archive entity {index} ({}) name must be non-blank",
358 entity.id
359 )));
360 }
361 }
362 for (index, edge) in archive.edges.iter().enumerate() {
363 let record = format!("edge[{index}]");
364 crate::operations::validate_edge_weight(edge.weight)?;
365 crate::secret_gate::reject_reserved_secret_gate_property(edge.properties.as_ref())?;
370 if let Some(p) = edge.properties.as_ref() {
371 crate::secret_gate::check_json_at(p, &record, "properties")?;
372 }
373 }
374
375 Ok(validation)
376 }
377
378 pub async fn import_kg_with_policy(
382 &self,
383 archive: &KgArchive,
384 token: &NamespaceToken,
385 policy: ImportKindPolicy,
386 ) -> RuntimeResult<ImportSummary> {
387 let validation = self.validate_kg_import(archive, policy)?;
388 for kind in validation.unknown_entity_kinds {
389 tracing::warn!(target: "khive.import.schema",
390 kind = %crate::secret_gate::bounded_masked_log_text(&kind),
391 "preserving unknown entity kind during explicit import");
392 }
393 let ns = token.namespace().as_str().to_owned();
394
395 let store = self.entities(token)?;
396 let mut entities_imported = 0usize;
397 let mut embedding_truncation = crate::retrieval::EmbeddingTruncationReport::default();
398 for page in archive
399 .entities
400 .chunks(crate::retrieval::EMBEDDING_BATCH_PAGE_SIZE)
401 {
402 let entities: Vec<khive_storage::entity::Entity> = page
403 .iter()
404 .map(|ee| khive_storage::entity::Entity {
405 entity_type: ee.entity_type.clone(),
406 description: ee.description.clone(),
407 properties: ee.properties.clone(),
408 tags: ee.tags.clone(),
409 ..khive_storage::entity::Entity::minimal(
410 ee.id,
411 ns.clone(),
412 ee.kind.clone(),
413 ee.name.clone(),
414 ee.created_at.timestamp_micros(),
415 ee.updated_at.timestamp_micros(),
416 )
417 })
418 .collect();
419 let texts: Vec<String> = entities
420 .iter()
421 .map(crate::curation::entity_embedding_text)
422 .collect();
423 let mut embeddings: Vec<HashMap<String, crate::retrieval::DocumentEmbeddingOutcome>> =
424 (0..entities.len()).map(|_| HashMap::new()).collect();
425 for model_name in self.registered_embedding_model_names() {
426 match self
427 .embed_document_batch_with_model_outcomes_for_token(token, &model_name, &texts)
428 .await
429 {
430 Ok(outcomes) => {
431 for (slot, outcome) in embeddings.iter_mut().zip(outcomes) {
432 slot.insert(model_name.clone(), outcome);
433 }
434 }
435 Err(error) => {
436 tracing::warn!(
437 model = %model_name,
438 error = %error,
439 "import_kg: batch embed failed; retrying records individually"
440 );
441 }
442 }
443 }
444 for (entity, outcomes) in entities.into_iter().zip(embeddings) {
445 store.upsert_entity(entity.clone()).await?;
446 embedding_truncation.merge(
449 self.reindex_entity_with_precomputed(token, &entity, outcomes)
450 .await?,
451 );
452 entities_imported += 1;
453 }
454 }
455
456 let graph = self.graph(token)?;
460 let mut edges_imported = 0usize;
461 let mut skipped_edges = Vec::new();
462 for ee in &archive.edges {
463 let source_ok = match self.get_entity(token, ee.source).await {
464 Ok(_) => true,
465 Err(RuntimeError::NotFound(_)) => false,
466 Err(e) => return Err(e),
467 };
468 if !source_ok {
469 tracing::warn!(
470 source = %ee.source,
471 target = %ee.target,
472 relation = ?ee.relation,
473 "import_kg: skipping edge — source entity not found in namespace {ns:?}"
474 );
475 skipped_edges.push(ImportSkippedEdge {
476 edge_id: ee.edge_id,
477 source: ee.source,
478 target: ee.target,
479 relation: ee.relation,
480 cause: ImportEdgeSkipCause::MissingSource,
481 error: None,
482 });
483 continue;
484 }
485 let target_ok = match self.get_entity(token, ee.target).await {
486 Ok(_) => true,
487 Err(RuntimeError::NotFound(_)) => false,
488 Err(e) => return Err(e),
489 };
490 if !target_ok {
491 tracing::warn!(
492 source = %ee.source,
493 target = %ee.target,
494 relation = ?ee.relation,
495 "import_kg: skipping edge — target entity not found in namespace {ns:?}"
496 );
497 skipped_edges.push(ImportSkippedEdge {
498 edge_id: ee.edge_id,
499 source: ee.source,
500 target: ee.target,
501 relation: ee.relation,
502 cause: ImportEdgeSkipCause::MissingTarget,
503 error: None,
504 });
505 continue;
506 }
507 match self
512 .validate_edge_relation_endpoints(token, ee.source, ee.target, ee.relation)
513 .await
514 {
515 Ok(_) => {}
516 Err(e @ (RuntimeError::InvalidInput(_) | RuntimeError::NotFound(_))) => {
517 tracing::warn!(
518 source = %ee.source,
519 target = %ee.target,
520 relation = ?ee.relation,
521 error = %e,
522 "import_kg: skipping edge — endpoint contract violation in namespace {ns:?}"
523 );
524 skipped_edges.push(ImportSkippedEdge {
525 edge_id: ee.edge_id,
526 source: ee.source,
527 target: ee.target,
528 relation: ee.relation,
529 cause: ImportEdgeSkipCause::EndpointContract,
530 error: Some(e.to_string()),
531 });
532 continue;
533 }
534 Err(e) => return Err(e),
535 }
536 let edge = khive_storage::types::Edge {
537 id: LinkId::from(ee.edge_id),
538 namespace: ns.clone(),
539 source_id: ee.source,
540 target_id: ee.target,
541 relation: ee.relation,
542 weight: ee.weight,
543 created_at: ee.created_at,
544 updated_at: ee.updated_at,
545 deleted_at: None,
546 metadata: ee.properties.clone(),
547 target_backend: None,
548 };
549 graph.upsert_edge(edge).await?;
550 edges_imported += 1;
551 }
552
553 Ok(ImportSummary {
554 entities_imported,
555 edges_imported,
556 edges_skipped: skipped_edges.len(),
557 skipped_edges,
558 embedding_truncation,
559 })
560 }
561
562 pub async fn import_kg_json(
564 &self,
565 json: &str,
566 token: &NamespaceToken,
567 ) -> RuntimeResult<ImportSummary> {
568 self.import_kg_json_with_policy(json, token, ImportKindPolicy::Strict)
569 .await
570 }
571
572 pub async fn import_kg_json_with_policy(
575 &self,
576 json: &str,
577 token: &NamespaceToken,
578 policy: ImportKindPolicy,
579 ) -> RuntimeResult<ImportSummary> {
580 let archive: KgArchive =
581 serde_json::from_str(json).map_err(|e| RuntimeError::InvalidInput(e.to_string()))?;
582 self.import_kg_with_policy(&archive, token, policy).await
583 }
584}
585
586#[cfg(test)]
587#[path = "portability_import_policy_tests.rs"]
588mod import_policy_tests;
589
590#[cfg(test)]
595mod tests {
596 use std::sync::atomic::{AtomicUsize, Ordering};
597 use std::sync::Arc;
598
599 use super::*;
600 use crate::runtime::{KhiveRuntime, NamespaceToken};
601 use crate::{EmbedderProvider, Namespace};
602 use async_trait::async_trait;
603 use khive_storage::EdgeRelation;
604 use lattice_embed::{EmbedError, EmbeddingModel, EmbeddingService, MAX_TEXT_BYTES};
605
606 const IMPORT_TEST_MODEL: &str = "all-minilm-l6-v2";
607 const IMPORT_BATCH_MODEL: &str = "import-batch-model";
608 const IMPORT_BATCH_MODEL_TWO: &str = "import-batch-model-two";
609
610 struct CountingImportService {
611 calls: Arc<AtomicUsize>,
612 reject_poison: bool,
613 }
614
615 #[async_trait]
616 impl EmbeddingService for CountingImportService {
617 async fn embed(
618 &self,
619 texts: &[String],
620 _model: EmbeddingModel,
621 ) -> std::result::Result<Vec<Vec<f32>>, EmbedError> {
622 self.calls.fetch_add(1, Ordering::SeqCst);
623 if self.reject_poison
624 && (texts.len() > 1 || texts.iter().any(|text| text.contains("poison")))
625 {
626 return Err(EmbedError::InferenceFailed("poison input".into()));
627 }
628 Ok(texts
629 .iter()
630 .map(|text| {
631 vec![
632 text.len() as f32,
633 text.bytes().map(u32::from).sum::<u32>() as f32,
634 text.as_bytes().first().copied().unwrap_or_default() as f32,
635 text.as_bytes().last().copied().unwrap_or_default() as f32,
636 ]
637 })
638 .collect())
639 }
640
641 fn supports_model(&self, _model: EmbeddingModel) -> bool {
642 true
643 }
644
645 fn name(&self) -> &'static str {
646 "import-batch-counting-service"
647 }
648 }
649
650 struct CountingImportProvider {
651 name: &'static str,
652 calls: Arc<AtomicUsize>,
653 reject_poison: bool,
654 }
655
656 #[async_trait]
657 impl EmbedderProvider for CountingImportProvider {
658 fn name(&self) -> &str {
659 self.name
660 }
661
662 fn dimensions(&self) -> usize {
663 4
664 }
665
666 async fn build(&self) -> std::result::Result<Arc<dyn EmbeddingService>, RuntimeError> {
667 Ok(Arc::new(CountingImportService {
668 calls: Arc::clone(&self.calls),
669 reject_poison: self.reject_poison,
670 }))
671 }
672 }
673
674 fn batch_archive(names: impl IntoIterator<Item = String>) -> KgArchive {
675 KgArchive {
676 format: "khive-kg".to_string(),
677 version: "0.1".to_string(),
678 namespace: "local".to_string(),
679 exported_at: Utc::now(),
680 entities: names
681 .into_iter()
682 .map(|name| ExportedEntity {
683 id: Uuid::new_v4(),
684 kind: "concept".to_string(),
685 entity_type: None,
686 name,
687 description: Some("batch body".to_string()),
688 properties: None,
689 tags: vec![],
690 created_at: Utc::now(),
691 updated_at: Utc::now(),
692 })
693 .collect(),
694 edges: vec![],
695 }
696 }
697
698 struct ImportEmbeddingService;
699
700 #[async_trait]
701 impl EmbeddingService for ImportEmbeddingService {
702 async fn embed(
703 &self,
704 texts: &[String],
705 _model: EmbeddingModel,
706 ) -> std::result::Result<Vec<Vec<f32>>, EmbedError> {
707 let dimensions = EmbeddingModel::AllMiniLmL6V2.dimensions();
708 Ok(texts.iter().map(|_| vec![1.0; dimensions]).collect())
709 }
710
711 fn supports_model(&self, _model: EmbeddingModel) -> bool {
712 true
713 }
714
715 fn name(&self) -> &'static str {
716 "kg-import-truncation-test"
717 }
718 }
719
720 struct ImportEmbeddingProvider;
721
722 #[async_trait]
723 impl EmbedderProvider for ImportEmbeddingProvider {
724 fn name(&self) -> &str {
725 IMPORT_TEST_MODEL
726 }
727
728 fn dimensions(&self) -> usize {
729 EmbeddingModel::AllMiniLmL6V2.dimensions()
730 }
731
732 async fn build(&self) -> std::result::Result<Arc<dyn EmbeddingService>, RuntimeError> {
733 Ok(Arc::new(ImportEmbeddingService))
734 }
735 }
736
737 async fn make_rt() -> KhiveRuntime {
738 KhiveRuntime::memory().expect("in-memory runtime")
739 }
740
741 fn make_rt_with_embedder() -> KhiveRuntime {
742 let runtime = KhiveRuntime::memory().expect("in-memory runtime");
743 runtime.register_embedder(ImportEmbeddingProvider);
744 runtime
745 }
746
747 #[test]
748 fn archive_timestamp_falls_back_only_when_absent_and_converts_offsets_to_utc() {
749 let fallback = DateTime::parse_from_rfc3339("2020-01-01T00:00:00Z")
750 .unwrap()
751 .with_timezone(&Utc);
752 assert_eq!(parse_archive_timestamp(None, fallback).unwrap(), fallback);
753
754 let utc = parse_archive_timestamp(Some("2024-05-06T07:08:09Z"), fallback).unwrap();
755 assert_eq!(utc.to_rfc3339(), "2024-05-06T07:08:09+00:00");
756
757 let offset = parse_archive_timestamp(Some("2024-05-06T09:08:09+02:00"), fallback).unwrap();
758 assert_eq!(offset, utc);
759
760 let err = parse_archive_timestamp(Some("2024-05-06 07:08:09"), fallback).unwrap_err();
761 let text = format!("{err:#}");
762 assert!(text.contains("\"2024-05-06 07:08:09\""), "{text}");
763 assert!(text.contains("must be RFC3339"), "{text}");
764 assert!(
765 err.chain().count() > 1,
766 "parser source error is kept: {text}"
767 );
768 }
769
770 #[tokio::test]
771 async fn import_batches_provider_calls_and_matches_per_record_reindex_vectors() {
772 let token = NamespaceToken::local();
773 let archive = batch_archive((0..257).map(|index| format!("Batch entity {index}")));
774 let ids: Vec<Uuid> = archive.entities.iter().map(|entity| entity.id).collect();
775 let calls = Arc::new(AtomicUsize::new(0));
776 let runtime = KhiveRuntime::memory().unwrap();
777 runtime.register_embedder(CountingImportProvider {
778 name: IMPORT_BATCH_MODEL,
779 calls: Arc::clone(&calls),
780 reject_poison: false,
781 });
782 let usage = crate::usage::UsageContext::new();
783 let summary = crate::usage::scope(usage.clone(), runtime.import_kg(&archive, &token))
784 .await
785 .unwrap();
786 assert_eq!(summary.entities_imported, ids.len());
787 assert_eq!(usage.snapshot()["embed_calls"], ids.len() as u64);
788 assert_eq!(
789 calls.load(Ordering::SeqCst),
790 2,
791 "257 records need two provider batches"
792 );
793
794 let singles = KhiveRuntime::memory().unwrap();
795 let singleton_calls = Arc::new(AtomicUsize::new(0));
796 singles.register_embedder(CountingImportProvider {
797 name: IMPORT_BATCH_MODEL,
798 calls: Arc::clone(&singleton_calls),
799 reject_poison: false,
800 });
801 for id in &ids {
802 let entity = runtime.get_entity(&token, *id).await.unwrap();
803 singles
804 .entities(&token)
805 .unwrap()
806 .upsert_entity(entity.clone())
807 .await
808 .unwrap();
809 singles.reindex_entity(&token, &entity).await.unwrap();
810 }
811 assert_eq!(singleton_calls.load(Ordering::SeqCst), ids.len());
812 let batched_vectors = runtime
813 .vectors_for_model(&token, IMPORT_BATCH_MODEL)
814 .unwrap()
815 .get_vectors(&ids, "local", "entity.body")
816 .await
817 .unwrap();
818 let singleton_vectors = singles
819 .vectors_for_model(&token, IMPORT_BATCH_MODEL)
820 .unwrap()
821 .get_vectors(&ids, "local", "entity.body")
822 .await
823 .unwrap();
824 assert_eq!(batched_vectors.len(), ids.len());
825 assert_eq!(batched_vectors, singleton_vectors);
826 }
827
828 #[tokio::test]
829 async fn import_batches_each_registered_model_once_per_page() {
830 let token = NamespaceToken::local();
831 let archive = batch_archive(["alpha", "beta", "gamma"].map(str::to_string));
832 let calls_a = Arc::new(AtomicUsize::new(0));
833 let calls_b = Arc::new(AtomicUsize::new(0));
834 let runtime = KhiveRuntime::memory().unwrap();
835 for (name, calls) in [
836 (IMPORT_BATCH_MODEL, Arc::clone(&calls_a)),
837 (IMPORT_BATCH_MODEL_TWO, Arc::clone(&calls_b)),
838 ] {
839 runtime.register_embedder(CountingImportProvider {
840 name,
841 calls,
842 reject_poison: false,
843 });
844 }
845 runtime.import_kg(&archive, &token).await.unwrap();
846 assert_eq!(calls_a.load(Ordering::SeqCst), 1);
847 assert_eq!(calls_b.load(Ordering::SeqCst), 1);
848 let ids: Vec<Uuid> = archive.entities.iter().map(|entity| entity.id).collect();
849 for model in [IMPORT_BATCH_MODEL, IMPORT_BATCH_MODEL_TWO] {
850 assert_eq!(
851 runtime
852 .vectors_for_model(&token, model)
853 .unwrap()
854 .get_vectors(&ids, "local", "entity.body")
855 .await
856 .unwrap()
857 .len(),
858 ids.len()
859 );
860 }
861 }
862
863 #[tokio::test]
864 async fn import_failed_batch_isolates_poison_without_duplicate_vector_writes() {
865 use khive_storage::types::{SqlStatement, SqlValue};
866
867 let token = NamespaceToken::local();
868 let archive = batch_archive(["good one", "poison", "good two"].map(str::to_string));
869 let ids: Vec<Uuid> = archive.entities.iter().map(|entity| entity.id).collect();
870 let calls = Arc::new(AtomicUsize::new(0));
871 let runtime = KhiveRuntime::memory().unwrap();
872 runtime.register_embedder(CountingImportProvider {
873 name: IMPORT_BATCH_MODEL,
874 calls: Arc::clone(&calls),
875 reject_poison: true,
876 });
877 let summary = runtime.import_kg(&archive, &token).await.unwrap();
878 assert_eq!(summary.entities_imported, 3);
879 assert_eq!(
880 calls.load(Ordering::SeqCst),
881 4,
882 "one failed page plus three singleton retries"
883 );
884 let vectors = runtime
885 .vectors_for_model(&token, IMPORT_BATCH_MODEL)
886 .unwrap()
887 .get_vectors(&ids, "local", "entity.body")
888 .await
889 .unwrap();
890 assert_eq!(vectors.len(), 2);
891 assert!(!vectors.contains_key(&ids[1]));
892 let mut reader = runtime.sql().reader().await.unwrap();
893 let rows = reader
894 .query_all(SqlStatement {
895 sql: "SELECT COUNT(*) FROM ann_write_log WHERE embedding_model = ?1 AND op = 'upsert'".into(),
896 params: vec![SqlValue::Text(IMPORT_BATCH_MODEL.into())],
897 label: Some("import-batch-upsert-count".into()),
898 })
899 .await
900 .unwrap();
901 assert!(matches!(&rows[0].columns[0].value, SqlValue::Integer(2)));
902 let provenance = reader
903 .query_scalar(SqlStatement {
904 sql: "SELECT COUNT(*) FROM vector_provenance WHERE model_key = ?1".into(),
905 params: vec![SqlValue::Text(crate::config::sanitize_key(
906 IMPORT_BATCH_MODEL,
907 ))],
908 label: Some("import-batch-provenance-count".into()),
909 })
910 .await
911 .unwrap();
912 assert!(matches!(provenance, Some(SqlValue::Integer(0))));
913 }
914
915 #[tokio::test]
916 async fn import_summary_reports_actual_embedding_truncation() {
917 let runtime = make_rt_with_embedder();
918 let token = NamespaceToken::local();
919 let archive = KgArchive {
920 format: "khive-kg".to_string(),
921 version: "0.1".to_string(),
922 namespace: "local".to_string(),
923 exported_at: Utc::now(),
924 entities: vec![ExportedEntity {
925 id: Uuid::new_v4(),
926 kind: "concept".to_string(),
927 entity_type: None,
928 name: "Imported long entity".to_string(),
929 description: Some("x".repeat(MAX_TEXT_BYTES + 1)),
930 properties: None,
931 tags: vec![],
932 created_at: Utc::now(),
933 updated_at: Utc::now(),
934 }],
935 edges: vec![],
936 };
937
938 let summary = runtime
939 .import_kg(&archive, &token)
940 .await
941 .expect("import with bounded embedding input");
942
943 assert_eq!(summary.entities_imported, 1);
944 assert_eq!(summary.embedding_truncation.truncated, 1);
945 assert!(summary.embedding_truncation.discarded_bytes > 0);
946 let wire = serde_json::to_value(&summary).expect("serialize import summary");
947 assert_eq!(wire["embedding_truncation"]["truncated"], 1);
948 }
949
950 #[tokio::test]
951 async fn import_rejects_whitespace_name_before_any_entity_write() {
952 let runtime = make_rt().await;
953 let token = NamespaceToken::local();
954 let valid_id = Uuid::new_v4();
955 let archive = KgArchive {
956 format: "khive-kg".to_string(),
957 version: "0.1".to_string(),
958 namespace: "local".to_string(),
959 exported_at: Utc::now(),
960 entities: vec![
961 ExportedEntity {
962 id: valid_id,
963 kind: "concept".to_string(),
964 entity_type: None,
965 name: "Must not be written".to_string(),
966 description: None,
967 properties: None,
968 tags: vec![],
969 created_at: Utc::now(),
970 updated_at: Utc::now(),
971 },
972 ExportedEntity {
973 id: Uuid::new_v4(),
974 kind: "concept".to_string(),
975 entity_type: None,
976 name: " \t\n ".to_string(),
977 description: None,
978 properties: None,
979 tags: vec![],
980 created_at: Utc::now(),
981 updated_at: Utc::now(),
982 },
983 ],
984 edges: vec![],
985 };
986
987 let err = runtime
988 .import_kg(&archive, &token)
989 .await
990 .expect_err("whitespace-only entity names must fail the whole import");
991 assert!(
992 err.to_string().contains("non-blank"),
993 "error must explain the name invariant: {err}"
994 );
995 assert!(
996 runtime.get_entity(&token, valid_id).await.is_err(),
997 "deterministic validation must finish before the first entity write"
998 );
999 }
1000
1001 #[tokio::test]
1003 async fn roundtrip_entities_and_edges() {
1004 let src = make_rt().await;
1005 let tok = NamespaceToken::local();
1006 let e1 = src
1007 .create_entity(
1008 &tok,
1009 "concept",
1010 None,
1011 "FlashAttention",
1012 Some("fast attention"),
1013 None,
1014 vec![],
1015 )
1016 .await
1017 .unwrap();
1018 let e2 = src
1019 .create_entity(
1020 &tok,
1021 "concept",
1022 None,
1023 "FlashAttention-2",
1024 None,
1025 None,
1026 vec![],
1027 )
1028 .await
1029 .unwrap();
1030 let e3 = src
1031 .create_entity(
1032 &tok,
1033 "person",
1034 None,
1035 "Tri Dao",
1036 None,
1037 None,
1038 vec!["author".into()],
1039 )
1040 .await
1041 .unwrap();
1042 src.link(&tok, e2.id, e1.id, EdgeRelation::Extends, 1.0, None)
1043 .await
1044 .unwrap();
1045 src.link(&tok, e1.id, e3.id, EdgeRelation::IntroducedBy, 0.9, None)
1046 .await
1047 .unwrap();
1048
1049 let archive = src.export_kg(&tok).await.unwrap();
1050 assert_eq!(archive.entities.len(), 3);
1051 assert_eq!(archive.edges.len(), 2);
1052 assert_eq!(archive.format, "khive-kg");
1053 assert_eq!(archive.version, "0.1");
1054
1055 let dst = make_rt().await;
1056 let summary = dst.import_kg(&archive, &tok).await.unwrap();
1057 assert_eq!(summary.entities_imported, 3);
1058 assert_eq!(summary.edges_imported, 2);
1059
1060 let got = dst.get_entity(&tok, e1.id).await.unwrap();
1061 assert_eq!(got.name, "FlashAttention");
1062 assert_eq!(got.description.as_deref(), Some("fast attention"));
1063 }
1064
1065 #[tokio::test]
1067 async fn json_roundtrip() {
1068 let src = make_rt().await;
1069 let tok = NamespaceToken::local();
1070 let e1 = src
1071 .create_entity(
1072 &tok,
1073 "concept",
1074 None,
1075 "LoRA",
1076 Some("low-rank adaptation"),
1077 Some(serde_json::json!({"year": "2021"})),
1078 vec!["fine-tuning".into()],
1079 )
1080 .await
1081 .unwrap();
1082 let e2 = src
1083 .create_entity(&tok, "concept", None, "QLoRA", None, None, vec![])
1084 .await
1085 .unwrap();
1086 src.link(&tok, e2.id, e1.id, EdgeRelation::VariantOf, 0.9, None)
1087 .await
1088 .unwrap();
1089
1090 let json_str = src.export_kg_json(&tok).await.unwrap();
1091 assert!(json_str.contains("khive-kg"));
1092
1093 let dst = make_rt().await;
1094 let summary = dst.import_kg_json(&json_str, &tok).await.unwrap();
1095 assert_eq!(summary.entities_imported, 2);
1096 assert_eq!(summary.edges_imported, 1);
1097
1098 let got = dst.get_entity(&tok, e1.id).await.unwrap();
1099 assert_eq!(got.tags, vec!["fine-tuning"]);
1100 }
1101
1102 #[tokio::test]
1109 async fn namespace_targeting() {
1110 let src = make_rt().await;
1111 let tok_a = NamespaceToken::for_namespace(Namespace::parse("a").unwrap());
1112 let tok_b = NamespaceToken::for_namespace(Namespace::parse("b").unwrap());
1113 src.create_entity(&tok_a, "concept", None, "Sinkhorn", None, None, vec![])
1114 .await
1115 .unwrap();
1116
1117 let archive = src.export_kg(&tok_a).await.unwrap();
1118 assert_eq!(archive.namespace, "a");
1119
1120 let dst = make_rt().await;
1121 let summary = dst.import_kg(&archive, &tok_b).await.unwrap();
1122 assert_eq!(summary.entities_imported, 1);
1123
1124 let in_b = dst.list_entities(&tok_b, None, None, 100, 0).await.unwrap();
1125 assert_eq!(in_b.len(), 1);
1126 assert_eq!(in_b[0].name, "Sinkhorn");
1127
1128 let in_a = src.list_entities(&tok_a, None, None, 100, 0).await.unwrap();
1129 assert_eq!(in_a.len(), 1);
1130
1131 let dst_a = dst.list_entities(&tok_a, None, None, 100, 0).await.unwrap();
1132 assert_eq!(dst_a.len(), 0);
1133 }
1134
1135 #[tokio::test]
1137 async fn format_validation_rejects_wrong_format() {
1138 let rt = make_rt().await;
1139 let tok = NamespaceToken::local();
1140 let bad = KgArchive {
1141 format: "wrong".to_string(),
1142 version: "0.1".to_string(),
1143 namespace: "local".to_string(),
1144 exported_at: Utc::now(),
1145 entities: vec![],
1146 edges: vec![],
1147 };
1148 let err = rt.import_kg(&bad, &tok).await.unwrap_err();
1149 assert!(matches!(err, RuntimeError::InvalidInput(_)));
1150 }
1151
1152 #[tokio::test]
1154 async fn import_unsupported_archive_version_returns_error() {
1155 let rt = make_rt().await;
1156 let tok = NamespaceToken::local();
1157 let bad = KgArchive {
1158 format: "khive-kg".to_string(),
1159 version: "999.0".to_string(),
1160 namespace: "local".to_string(),
1161 exported_at: Utc::now(),
1162 entities: vec![],
1163 edges: vec![],
1164 };
1165 let err = rt.import_kg(&bad, &tok).await.unwrap_err();
1166 assert!(
1167 matches!(err, RuntimeError::InvalidInput(_)),
1168 "expected InvalidInput, got {err:?}"
1169 );
1170 if let RuntimeError::InvalidInput(msg) = err {
1171 assert!(
1172 msg.contains("999.0"),
1173 "error message should mention the unsupported version, got: {msg:?}"
1174 );
1175 }
1176 }
1177
1178 #[tokio::test]
1182 async fn import_entity_with_reserved_secret_gate_property_is_rejected() {
1183 let rt = make_rt().await;
1184 let tok = NamespaceToken::local();
1185 let archive = KgArchive {
1186 format: "khive-kg".to_string(),
1187 version: "0.1".to_string(),
1188 namespace: "local".to_string(),
1189 exported_at: Utc::now(),
1190 entities: vec![ExportedEntity {
1191 id: Uuid::new_v4(),
1192 kind: "concept".to_string(),
1193 entity_type: None,
1194 name: "ReservedKeyImport".to_string(),
1195 description: None,
1196 properties: Some(serde_json::json!({
1197 "khive:secret_gate": "exempted:content-sha256-manifest-v1"
1198 })),
1199 tags: vec![],
1200 created_at: Utc::now(),
1201 updated_at: Utc::now(),
1202 }],
1203 edges: vec![],
1204 };
1205
1206 let err = rt.import_kg(&archive, &tok).await.unwrap_err();
1207 assert!(
1208 matches!(err, RuntimeError::InvalidInput(ref msg) if msg.contains("khive:secret_gate") && msg.contains("runtime-owned")),
1209 "expected a reservation rejection, got {err:?}"
1210 );
1211
1212 let entities = rt.list_entities(&tok, None, None, 100, 0).await.unwrap();
1213 assert!(
1214 !entities.iter().any(|e| e.name == "ReservedKeyImport"),
1215 "rejected archive import must not create the entity"
1216 );
1217 }
1218
1219 #[tokio::test]
1224 async fn import_edge_with_reserved_secret_gate_property_is_rejected() {
1225 let rt = make_rt().await;
1226 let tok = NamespaceToken::local();
1227 let a = Uuid::new_v4();
1228 let b = Uuid::new_v4();
1229 let mk = |id: Uuid, name: &str| ExportedEntity {
1230 id,
1231 kind: "concept".to_string(),
1232 entity_type: None,
1233 name: name.to_string(),
1234 description: None,
1235 properties: None,
1236 tags: vec![],
1237 created_at: Utc::now(),
1238 updated_at: Utc::now(),
1239 };
1240 let archive = KgArchive {
1241 format: "khive-kg".to_string(),
1242 version: "0.1".to_string(),
1243 namespace: "local".to_string(),
1244 exported_at: Utc::now(),
1245 entities: vec![mk(a, "EdgeGateA"), mk(b, "EdgeGateB")],
1246 edges: vec![ExportedEdge {
1247 edge_id: Uuid::new_v4(),
1248 source: a,
1249 target: b,
1250 relation: EdgeRelation::Extends,
1251 weight: 0.5,
1252 properties: Some(serde_json::json!({
1253 "khive:secret_gate": "exempted:content-sha256-manifest-v1"
1254 })),
1255 created_at: Utc::now(),
1256 updated_at: Utc::now(),
1257 }],
1258 };
1259
1260 let err = rt.import_kg(&archive, &tok).await.unwrap_err();
1261 assert!(
1262 matches!(err, RuntimeError::InvalidInput(ref msg) if msg.contains("khive:secret_gate") && msg.contains("runtime-owned")),
1263 "expected a reservation rejection, got {err:?}"
1264 );
1265
1266 let entities = rt.list_entities(&tok, None, None, 100, 0).await.unwrap();
1269 assert!(
1270 !entities.iter().any(|e| e.name == "EdgeGateA"),
1271 "rejected archive import must not create any records"
1272 );
1273 }
1274
1275 #[tokio::test]
1279 async fn import_edge_with_credential_shaped_property_is_rejected() {
1280 let rt = make_rt().await;
1281 let tok = NamespaceToken::local();
1282 let a = Uuid::new_v4();
1283 let b = Uuid::new_v4();
1284 let mk = |id: Uuid, name: &str| ExportedEntity {
1285 id,
1286 kind: "concept".to_string(),
1287 entity_type: None,
1288 name: name.to_string(),
1289 description: None,
1290 properties: None,
1291 tags: vec![],
1292 created_at: Utc::now(),
1293 updated_at: Utc::now(),
1294 };
1295 let archive = KgArchive {
1296 format: "khive-kg".to_string(),
1297 version: "0.1".to_string(),
1298 namespace: "local".to_string(),
1299 exported_at: Utc::now(),
1300 entities: vec![mk(a, "EdgeScanA"), mk(b, "EdgeScanB")],
1301 edges: vec![ExportedEdge {
1302 edge_id: Uuid::new_v4(),
1303 source: a,
1304 target: b,
1305 relation: EdgeRelation::Extends,
1306 weight: 0.5,
1307 properties: Some(serde_json::json!({
1308 "api_key": "AKIAFAKEKEY1234567890"
1309 })),
1310 created_at: Utc::now(),
1311 updated_at: Utc::now(),
1312 }],
1313 };
1314
1315 let err = rt.import_kg(&archive, &tok).await.unwrap_err();
1316 assert!(
1317 matches!(err, RuntimeError::SecretDetected(_)),
1318 "expected the blocking scanner to reject the value, got {err:?}"
1319 );
1320
1321 let entities = rt.list_entities(&tok, None, None, 100, 0).await.unwrap();
1322 assert!(
1323 !entities.iter().any(|e| e.name == "EdgeScanA"),
1324 "rejected archive import must not create any records"
1325 );
1326 }
1327
1328 #[test]
1330 fn invalid_relation_rejected_at_deserialize() {
1331 let json = r#"{
1332 "format":"khive-kg","version":"0.1","namespace":"local",
1333 "exported_at":"2026-01-01T00:00:00Z",
1334 "entities":[],
1335 "edges":[{"edge_id":"00000000-0000-0000-0000-000000000099",
1336 "source":"00000000-0000-0000-0000-000000000001",
1337 "target":"00000000-0000-0000-0000-000000000002",
1338 "relation":"related_to","weight":0.5}]
1339 }"#;
1340 let result: Result<KgArchive, _> = serde_json::from_str(json);
1341 assert!(
1342 result.is_err(),
1343 "non-canonical relation should fail to deserialize"
1344 );
1345 }
1346
1347 #[tokio::test]
1354 async fn import_edge_with_dangling_source_is_skipped() {
1355 let phantom_source = Uuid::parse_str("deadbeef-dead-4ead-dead-deadbeefcafe").unwrap();
1356
1357 let rt = make_rt().await;
1358 let tok = NamespaceToken::local();
1359 let real = rt
1360 .create_entity(&tok, "concept", None, "Real", None, None, vec![])
1361 .await
1362 .unwrap();
1363
1364 let archive = KgArchive {
1365 format: "khive-kg".to_string(),
1366 version: "0.1".to_string(),
1367 namespace: "local".to_string(),
1368 exported_at: Utc::now(),
1369 entities: vec![ExportedEntity {
1370 id: real.id,
1371 kind: "concept".to_string(),
1372 entity_type: None,
1373 name: "Real".to_string(),
1374 description: None,
1375 properties: None,
1376 tags: vec![],
1377 created_at: Utc::now(),
1378 updated_at: Utc::now(),
1379 }],
1380 edges: vec![ExportedEdge {
1381 edge_id: Uuid::new_v4(),
1382 source: phantom_source,
1383 target: real.id,
1384 relation: EdgeRelation::Extends,
1385 weight: 1.0,
1386 properties: None,
1387 created_at: Utc::now(),
1388 updated_at: Utc::now(),
1389 }],
1390 };
1391
1392 let dst = make_rt().await;
1393 let summary = dst.import_kg(&archive, &tok).await.unwrap();
1394 assert_eq!(summary.entities_imported, 1);
1395 assert_eq!(
1396 summary.edges_imported, 0,
1397 "dangling source must not be imported"
1398 );
1399 assert_eq!(
1400 summary.edges_skipped, 1,
1401 "dangling source must be counted as skipped"
1402 );
1403 }
1404
1405 #[tokio::test]
1410 async fn import_edge_with_dangling_target_is_skipped() {
1411 let phantom_target = Uuid::parse_str("cafebabe-cafe-4abe-cafe-cafebabecafe").unwrap();
1412
1413 let rt = make_rt().await;
1414 let tok = NamespaceToken::local();
1415 let real = rt
1416 .create_entity(&tok, "concept", None, "Source", None, None, vec![])
1417 .await
1418 .unwrap();
1419
1420 let archive = KgArchive {
1421 format: "khive-kg".to_string(),
1422 version: "0.1".to_string(),
1423 namespace: "local".to_string(),
1424 exported_at: Utc::now(),
1425 entities: vec![ExportedEntity {
1426 id: real.id,
1427 kind: "concept".to_string(),
1428 entity_type: None,
1429 name: "Source".to_string(),
1430 description: None,
1431 properties: None,
1432 tags: vec![],
1433 created_at: Utc::now(),
1434 updated_at: Utc::now(),
1435 }],
1436 edges: vec![ExportedEdge {
1437 edge_id: Uuid::new_v4(),
1438 source: real.id,
1439 target: phantom_target,
1440 relation: EdgeRelation::DependsOn,
1441 weight: 0.8,
1442 properties: None,
1443 created_at: Utc::now(),
1444 updated_at: Utc::now(),
1445 }],
1446 };
1447
1448 let dst = make_rt().await;
1449 let summary = dst.import_kg(&archive, &tok).await.unwrap();
1450 assert_eq!(summary.entities_imported, 1);
1451 assert_eq!(
1452 summary.edges_imported, 0,
1453 "dangling target must not be imported"
1454 );
1455 assert_eq!(
1456 summary.edges_skipped, 1,
1457 "dangling target must be counted as skipped"
1458 );
1459 }
1460
1461 #[tokio::test]
1466 async fn import_mixed_edges_reports_correct_counts() {
1467 let phantom = Uuid::parse_str("11111111-1111-4111-8111-111111111111").unwrap();
1468
1469 let src = make_rt().await;
1470 let tok = NamespaceToken::local();
1471 let a = src
1472 .create_entity(&tok, "concept", None, "A", None, None, vec![])
1473 .await
1474 .unwrap();
1475 let b = src
1476 .create_entity(&tok, "concept", None, "B", None, None, vec![])
1477 .await
1478 .unwrap();
1479 let c = src
1480 .create_entity(&tok, "concept", None, "C", None, None, vec![])
1481 .await
1482 .unwrap();
1483
1484 let archive = KgArchive {
1485 format: "khive-kg".to_string(),
1486 version: "0.1".to_string(),
1487 namespace: "local".to_string(),
1488 exported_at: Utc::now(),
1489 entities: vec![
1490 ExportedEntity {
1491 id: a.id,
1492 kind: "concept".to_string(),
1493 entity_type: None,
1494 name: "A".to_string(),
1495 description: None,
1496 properties: None,
1497 tags: vec![],
1498 created_at: Utc::now(),
1499 updated_at: Utc::now(),
1500 },
1501 ExportedEntity {
1502 id: b.id,
1503 kind: "concept".to_string(),
1504 entity_type: None,
1505 name: "B".to_string(),
1506 description: None,
1507 properties: None,
1508 tags: vec![],
1509 created_at: Utc::now(),
1510 updated_at: Utc::now(),
1511 },
1512 ExportedEntity {
1513 id: c.id,
1514 kind: "concept".to_string(),
1515 entity_type: None,
1516 name: "C".to_string(),
1517 description: None,
1518 properties: None,
1519 tags: vec![],
1520 created_at: Utc::now(),
1521 updated_at: Utc::now(),
1522 },
1523 ],
1524 edges: vec![
1525 ExportedEdge {
1526 edge_id: Uuid::new_v4(),
1527 source: a.id,
1528 target: b.id,
1529 relation: EdgeRelation::Extends,
1530 weight: 1.0,
1531 properties: None,
1532 created_at: Utc::now(),
1533 updated_at: Utc::now(),
1534 },
1535 ExportedEdge {
1536 edge_id: Uuid::new_v4(),
1537 source: b.id,
1538 target: c.id,
1539 relation: EdgeRelation::VariantOf,
1544 weight: 0.9,
1545 properties: None,
1546 created_at: Utc::now(),
1547 updated_at: Utc::now(),
1548 },
1549 ExportedEdge {
1550 edge_id: Uuid::new_v4(),
1551 source: a.id,
1552 target: phantom,
1553 relation: EdgeRelation::Enables,
1554 weight: 0.5,
1555 properties: None,
1556 created_at: Utc::now(),
1557 updated_at: Utc::now(),
1558 },
1559 ],
1560 };
1561
1562 let dst = make_rt().await;
1563 let summary = dst.import_kg(&archive, &tok).await.unwrap();
1564 assert_eq!(summary.entities_imported, 3);
1565 assert_eq!(
1566 summary.edges_imported, 2,
1567 "only valid edges must be imported"
1568 );
1569 assert_eq!(
1570 summary.edges_skipped, 1,
1571 "one dangling edge must be reported"
1572 );
1573 }
1574
1575 #[tokio::test]
1577 async fn import_all_valid_edges_reports_zero_skipped() {
1578 let src = make_rt().await;
1579 let tok = NamespaceToken::local();
1580 let e1 = src
1581 .create_entity(&tok, "concept", None, "E1", None, None, vec![])
1582 .await
1583 .unwrap();
1584 let e2 = src
1585 .create_entity(&tok, "concept", None, "E2", None, None, vec![])
1586 .await
1587 .unwrap();
1588 src.link(&tok, e1.id, e2.id, EdgeRelation::VariantOf, 0.7, None)
1589 .await
1590 .unwrap();
1591
1592 let archive = src.export_kg(&tok).await.unwrap();
1593 let dst = make_rt().await;
1594 let summary = dst.import_kg(&archive, &tok).await.unwrap();
1595 assert_eq!(summary.edges_imported, 1);
1596 assert_eq!(
1597 summary.edges_skipped, 0,
1598 "no edges should be skipped when all endpoints exist"
1599 );
1600 }
1601
1602 #[tokio::test]
1608 async fn import_edge_violating_endpoint_contract_is_skipped() {
1609 let src = make_rt().await;
1610 let tok = NamespaceToken::local();
1611 let e1 = src
1612 .create_entity(&tok, "concept", None, "E1", None, None, vec![])
1613 .await
1614 .unwrap();
1615 let e2 = src
1616 .create_entity(&tok, "concept", None, "E2", None, None, vec![])
1617 .await
1618 .unwrap();
1619
1620 let archive = KgArchive {
1621 format: "khive-kg".to_string(),
1622 version: "0.1".to_string(),
1623 namespace: "local".to_string(),
1624 exported_at: Utc::now(),
1625 entities: vec![
1626 ExportedEntity {
1627 id: e1.id,
1628 kind: "concept".to_string(),
1629 entity_type: None,
1630 name: "E1".to_string(),
1631 description: None,
1632 properties: None,
1633 tags: vec![],
1634 created_at: Utc::now(),
1635 updated_at: Utc::now(),
1636 },
1637 ExportedEntity {
1638 id: e2.id,
1639 kind: "concept".to_string(),
1640 entity_type: None,
1641 name: "E2".to_string(),
1642 description: None,
1643 properties: None,
1644 tags: vec![],
1645 created_at: Utc::now(),
1646 updated_at: Utc::now(),
1647 },
1648 ],
1649 edges: vec![ExportedEdge {
1650 edge_id: Uuid::new_v4(),
1651 source: e1.id,
1652 target: e2.id,
1653 relation: EdgeRelation::Precedes,
1654 weight: 1.0,
1655 properties: None,
1656 created_at: Utc::now(),
1657 updated_at: Utc::now(),
1658 }],
1659 };
1660
1661 let dst = make_rt().await;
1662 let summary = dst.import_kg(&archive, &tok).await.unwrap();
1663 assert_eq!(summary.entities_imported, 2);
1664 assert_eq!(
1665 summary.edges_imported, 0,
1666 "an endpoint-contract-violating edge must not be imported"
1667 );
1668 assert_eq!(
1669 summary.edges_skipped, 1,
1670 "an endpoint-contract-violating edge must be counted as skipped"
1671 );
1672 assert!(
1673 dst.neighbors(
1674 &tok,
1675 e1.id,
1676 khive_storage::types::Direction::Out,
1677 None,
1678 None
1679 )
1680 .await
1681 .unwrap()
1682 .is_empty(),
1683 "the contract-violating edge must not exist in the destination graph"
1684 );
1685 }
1686
1687 #[tokio::test]
1691 async fn export_kg_preserves_edge_id() {
1692 let rt = make_rt().await;
1693 let tok = NamespaceToken::local();
1694 let a = rt
1695 .create_entity(&tok, "concept", None, "Alpha", None, None, vec![])
1696 .await
1697 .unwrap();
1698 let b = rt
1699 .create_entity(&tok, "concept", None, "Beta", None, None, vec![])
1700 .await
1701 .unwrap();
1702 let stored_edge = rt
1703 .link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
1704 .await
1705 .unwrap();
1706 let stored_id: Uuid = stored_edge.id.into();
1707
1708 let archive = rt.export_kg(&tok).await.unwrap();
1709 assert_eq!(archive.edges.len(), 1);
1710 assert_eq!(
1711 archive.edges[0].edge_id, stored_id,
1712 "exported edge_id must equal the LinkId returned by link"
1713 );
1714 }
1715
1716 #[tokio::test]
1718 async fn import_kg_persists_edge_id_and_timestamps() {
1719 let src = make_rt().await;
1720 let tok = NamespaceToken::local();
1721 let a = src
1722 .create_entity(&tok, "concept", None, "Alpha", None, None, vec![])
1723 .await
1724 .unwrap();
1725 let b = src
1726 .create_entity(&tok, "concept", None, "Beta", None, None, vec![])
1727 .await
1728 .unwrap();
1729 let stored_edge = src
1730 .link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
1731 .await
1732 .unwrap();
1733 let original_id: Uuid = stored_edge.id.into();
1734
1735 let expected_created = chrono::DateTime::parse_from_rfc3339("2026-03-03T00:00:00Z")
1736 .unwrap()
1737 .with_timezone(&Utc);
1738 let expected_updated = chrono::DateTime::parse_from_rfc3339("2026-04-04T00:00:00Z")
1739 .unwrap()
1740 .with_timezone(&Utc);
1741 let mut archive = src.export_kg(&tok).await.unwrap();
1742 archive.edges[0].created_at = expected_created;
1743 archive.edges[0].updated_at = expected_updated;
1744 archive.edges[0].properties = Some(serde_json::json!({"confidence": 0.95}));
1745 let dst = make_rt().await;
1746 dst.import_kg(&archive, &tok).await.unwrap();
1747
1748 let imported_edge = dst.get_edge(&tok, original_id).await.unwrap();
1749 assert!(
1750 imported_edge.is_some(),
1751 "imported edge must be retrievable by the original edge_id"
1752 );
1753 let imported_edge = imported_edge.unwrap();
1754 assert_eq!(
1755 Uuid::from(imported_edge.id),
1756 original_id,
1757 "stored edge id must equal the archive edge_id"
1758 );
1759 assert_eq!(
1760 imported_edge.created_at, expected_created,
1761 "present edge created_at must not be replaced with import time"
1762 );
1763 assert_eq!(
1764 imported_edge.updated_at, expected_updated,
1765 "present edge updated_at must not be replaced with created_at or import time"
1766 );
1767 assert_eq!(
1768 imported_edge.metadata,
1769 Some(serde_json::json!({"confidence": 0.95})),
1770 "edge properties must persist as storage metadata"
1771 );
1772
1773 let reexported = dst.export_kg(&tok).await.unwrap();
1774 assert_eq!(
1775 reexported.edges[0].properties,
1776 Some(serde_json::json!({"confidence": 0.95})),
1777 "edge metadata must re-export as portable properties"
1778 );
1779 }
1780
1781 #[tokio::test]
1786 async fn old_archive_missing_edge_id_round_trips() {
1787 let src_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
1788 let tgt_id = Uuid::parse_str("00000000-0000-0000-0000-000000000002").unwrap();
1789
1790 let json = format!(
1792 r#"{{
1793 "format": "khive-kg",
1794 "version": "0.1",
1795 "namespace": "local",
1796 "exported_at": "2026-01-01T00:00:00Z",
1797 "entities": [
1798 {{"id":"{src_id}","kind":"concept","name":"SrcNode","created_at":"2026-01-01T00:00:00Z","updated_at":"2026-01-01T00:00:00Z"}},
1799 {{"id":"{tgt_id}","kind":"concept","name":"TgtNode","created_at":"2026-01-01T00:00:00Z","updated_at":"2026-01-01T00:00:00Z"}}
1800 ],
1801 "edges": [
1802 {{
1803 "source": "{src_id}",
1804 "target": "{tgt_id}",
1805 "relation": "extends",
1806 "weight": 0.9
1807 }}
1808 ]
1809 }}"#
1810 );
1811
1812 let archive: KgArchive = serde_json::from_str(&json)
1814 .expect("old archive without edge_id must deserialize successfully");
1815 assert_eq!(archive.edges.len(), 1);
1816 let generated_id = archive.edges[0].edge_id;
1817 assert_ne!(
1818 generated_id,
1819 Uuid::nil(),
1820 "missing edge_id in old archive must get a fresh non-nil UUID"
1821 );
1822
1823 let rt = make_rt().await;
1824 let tok = NamespaceToken::local();
1825 let summary = rt.import_kg(&archive, &tok).await.unwrap();
1826 assert_eq!(summary.entities_imported, 2);
1827 assert_eq!(
1828 summary.edges_imported, 1,
1829 "edge must be imported when both endpoints exist"
1830 );
1831
1832 let stored = rt.get_edge(&tok, generated_id).await.unwrap();
1833 assert!(
1834 stored.is_some(),
1835 "imported edge must be retrievable by the generated edge_id"
1836 );
1837 assert_eq!(
1838 Uuid::from(stored.unwrap().id),
1839 generated_id,
1840 "stored edge id must equal the generated edge_id"
1841 );
1842
1843 let re_archive = rt.export_kg(&tok).await.unwrap();
1844 assert_eq!(re_archive.edges.len(), 1);
1845 assert_eq!(
1846 re_archive.edges[0].edge_id, generated_id,
1847 "re-exported edge_id must equal the ID generated on first import"
1848 );
1849 }
1850
1851 #[tokio::test]
1856 async fn export_import_export_edge_id_equality() {
1857 let src = make_rt().await;
1858 let tok = NamespaceToken::local();
1859 let a = src
1860 .create_entity(&tok, "concept", None, "NodeA", None, None, vec![])
1861 .await
1862 .unwrap();
1863 let b = src
1864 .create_entity(&tok, "concept", None, "NodeB", None, None, vec![])
1865 .await
1866 .unwrap();
1867 let stored = src
1868 .link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
1869 .await
1870 .unwrap();
1871 let original_edge_id: Uuid = stored.id.into();
1872
1873 let archive1 = src.export_kg(&tok).await.unwrap();
1874 assert_eq!(archive1.edges.len(), 1);
1875 assert_eq!(
1876 archive1.edges[0].edge_id, original_edge_id,
1877 "first export must carry the stored edge_id"
1878 );
1879
1880 let dst = make_rt().await;
1881 dst.import_kg(&archive1, &tok).await.unwrap();
1882
1883 let archive2 = dst.export_kg(&tok).await.unwrap();
1884 assert_eq!(archive2.edges.len(), 1);
1885
1886 let re_edge = archive2
1887 .edges
1888 .iter()
1889 .find(|e| e.source == a.id && e.target == b.id && e.relation == EdgeRelation::Extends)
1890 .expect(
1891 "re-exported archive must contain the original edge by (source,target,relation)",
1892 );
1893 assert_eq!(
1894 re_edge.edge_id, original_edge_id,
1895 "edge_id must be identical across export → import → export"
1896 );
1897 }
1898
1899 fn skip_report_entity(id: Uuid, name: &str) -> ExportedEntity {
1901 let timestamp = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
1902 .unwrap()
1903 .with_timezone(&Utc);
1904 ExportedEntity {
1905 id,
1906 kind: "document".into(),
1907 entity_type: None,
1908 name: name.into(),
1909 description: None,
1910 properties: None,
1911 tags: vec![],
1912 created_at: timestamp,
1913 updated_at: timestamp,
1914 }
1915 }
1916
1917 fn skip_report_edge(
1918 id: u128,
1919 source: Uuid,
1920 target: Uuid,
1921 relation: EdgeRelation,
1922 ) -> ExportedEdge {
1923 let timestamp = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
1924 .unwrap()
1925 .with_timezone(&Utc);
1926 ExportedEdge {
1927 edge_id: Uuid::from_u128(id),
1928 source,
1929 target,
1930 relation,
1931 weight: 0.7,
1932 properties: None,
1933 created_at: timestamp,
1934 updated_at: timestamp,
1935 }
1936 }
1937
1938 fn skip_report_archive(entities: Vec<ExportedEntity>, edges: Vec<ExportedEdge>) -> KgArchive {
1939 KgArchive {
1940 format: "khive-kg".into(),
1941 version: "0.1".into(),
1942 namespace: "local".into(),
1943 exported_at: Utc::now(),
1944 entities,
1945 edges,
1946 }
1947 }
1948
1949 #[tokio::test]
1950 async fn skipped_edge_details_cover_three_causes_and_retain_the_valid_edge() {
1951 let runtime = make_rt().await;
1952 let token = NamespaceToken::local();
1953 let source = Uuid::from_u128(1);
1954 let target = Uuid::from_u128(2);
1955 let absent = Uuid::from_u128(3);
1956 let mut valid = skip_report_edge(101, source, target, EdgeRelation::LinksTo);
1957 valid.properties = Some(serde_json::json!({"retained": true}));
1958 let archive = skip_report_archive(
1959 vec![
1960 skip_report_entity(source, "Source"),
1961 skip_report_entity(target, "Target"),
1962 ],
1963 vec![
1964 skip_report_edge(100, absent, source, EdgeRelation::LinksTo),
1965 valid.clone(),
1966 skip_report_edge(102, source, absent, EdgeRelation::DependsOn),
1967 skip_report_edge(103, source, target, EdgeRelation::Annotates),
1968 ],
1969 );
1970 let summary = runtime.import_kg(&archive, &token).await.unwrap();
1971 let expected = vec![
1972 ImportSkippedEdge {
1973 edge_id: Uuid::from_u128(100),
1974 source: absent,
1975 target: source,
1976 relation: EdgeRelation::LinksTo,
1977 cause: ImportEdgeSkipCause::MissingSource,
1978 error: None,
1979 },
1980 ImportSkippedEdge {
1981 edge_id: Uuid::from_u128(102),
1982 source,
1983 target: absent,
1984 relation: EdgeRelation::DependsOn,
1985 cause: ImportEdgeSkipCause::MissingTarget,
1986 error: None,
1987 },
1988 ImportSkippedEdge {
1989 edge_id: Uuid::from_u128(103),
1990 source,
1991 target,
1992 relation: EdgeRelation::Annotates,
1993 cause: ImportEdgeSkipCause::EndpointContract,
1994 error: Some(format!(
1995 "invalid input: annotates source {source} must be a note"
1996 )),
1997 },
1998 ];
1999 assert_eq!(summary.entities_imported, 2);
2000 assert_eq!(summary.edges_imported, 1);
2001 assert_eq!(summary.edges_skipped, 3);
2002 assert_eq!(summary.skipped_edges, expected);
2003 let wire = serde_json::to_value(&summary).unwrap();
2004 assert_eq!(
2005 wire["skipped_edges"],
2006 serde_json::json!([
2007 {"edge_id":Uuid::from_u128(100),"source":absent,"target":source,
2008 "relation":"links_to","cause":"missing_source"},
2009 {"edge_id":Uuid::from_u128(102),"source":source,"target":absent,
2010 "relation":"depends_on","cause":"missing_target"},
2011 {"edge_id":Uuid::from_u128(103),"source":source,"target":target,
2012 "relation":"annotates","cause":"endpoint_contract",
2013 "error":format!("invalid input: annotates source {source} must be a note")}
2014 ])
2015 );
2016 let decoded: ImportSummary = serde_json::from_value(wire).unwrap();
2017 assert_eq!(decoded.skipped_edges, expected);
2018 assert_eq!(decoded.edges_skipped, 3);
2019 for id in [100, 102, 103] {
2020 assert!(runtime
2021 .get_edge(&token, Uuid::from_u128(id))
2022 .await
2023 .unwrap()
2024 .is_none());
2025 }
2026 let stored = runtime
2027 .get_edge(&token, valid.edge_id)
2028 .await
2029 .unwrap()
2030 .unwrap();
2031 assert_eq!(Uuid::from(stored.id), valid.edge_id);
2032 assert_eq!(stored.source_id, source);
2033 assert_eq!(stored.target_id, target);
2034 assert_eq!(stored.relation, EdgeRelation::LinksTo);
2035 assert_eq!(stored.metadata, valid.properties);
2036 assert_eq!(stored.created_at, valid.created_at);
2037 assert_eq!(stored.updated_at, valid.updated_at);
2038 let exported = runtime.export_kg(&token).await.unwrap();
2039 assert_eq!(exported.edges.len(), 1);
2040 assert_eq!(exported.edges[0].edge_id, valid.edge_id);
2041 }
2042
2043 #[tokio::test]
2044 async fn skipped_edge_details_preserve_duplicates_order_and_both_missing_precedence() {
2045 let runtime = make_rt().await;
2046 let token = NamespaceToken::local();
2047 let present = Uuid::from_u128(1);
2048 let absent = Uuid::from_u128(3);
2049 let other_absent = Uuid::from_u128(4);
2050 let repeated = skip_report_edge(111, absent, present, EdgeRelation::LinksTo);
2051 let archive = skip_report_archive(
2052 vec![skip_report_entity(present, "Present")],
2053 vec![
2054 skip_report_edge(110, present, absent, EdgeRelation::DependsOn),
2055 repeated.clone(),
2056 skip_report_edge(112, absent, present, EdgeRelation::LinksTo),
2057 repeated,
2058 skip_report_edge(113, absent, other_absent, EdgeRelation::Extends),
2059 ],
2060 );
2061 let summary = runtime.import_kg(&archive, &token).await.unwrap();
2062 assert_eq!(summary.edges_imported, 0);
2063 assert_eq!(summary.edges_skipped, 5);
2064 assert_eq!(
2065 summary.skipped_edges,
2066 vec![
2067 ImportSkippedEdge {
2068 edge_id: Uuid::from_u128(110),
2069 source: present,
2070 target: absent,
2071 relation: EdgeRelation::DependsOn,
2072 cause: ImportEdgeSkipCause::MissingTarget,
2073 error: None,
2074 },
2075 ImportSkippedEdge {
2076 edge_id: Uuid::from_u128(111),
2077 source: absent,
2078 target: present,
2079 relation: EdgeRelation::LinksTo,
2080 cause: ImportEdgeSkipCause::MissingSource,
2081 error: None,
2082 },
2083 ImportSkippedEdge {
2084 edge_id: Uuid::from_u128(112),
2085 source: absent,
2086 target: present,
2087 relation: EdgeRelation::LinksTo,
2088 cause: ImportEdgeSkipCause::MissingSource,
2089 error: None,
2090 },
2091 ImportSkippedEdge {
2092 edge_id: Uuid::from_u128(111),
2093 source: absent,
2094 target: present,
2095 relation: EdgeRelation::LinksTo,
2096 cause: ImportEdgeSkipCause::MissingSource,
2097 error: None,
2098 },
2099 ImportSkippedEdge {
2100 edge_id: Uuid::from_u128(113),
2101 source: absent,
2102 target: other_absent,
2103 relation: EdgeRelation::Extends,
2104 cause: ImportEdgeSkipCause::MissingSource,
2105 error: None,
2106 },
2107 ]
2108 );
2109 assert!(runtime.export_kg(&token).await.unwrap().edges.is_empty());
2110 }
2111
2112 #[tokio::test]
2113 async fn successful_imports_without_skips_keep_the_previous_json_shape() {
2114 let token = NamespaceToken::local();
2115 let a = Uuid::from_u128(1);
2116 let b = Uuid::from_u128(2);
2117 for (archive, entities, edges) in [
2118 (skip_report_archive(vec![], vec![]), 0, 0),
2119 (
2120 skip_report_archive(
2121 vec![skip_report_entity(a, "A"), skip_report_entity(b, "B")],
2122 vec![skip_report_edge(101, a, b, EdgeRelation::LinksTo)],
2123 ),
2124 2,
2125 1,
2126 ),
2127 ] {
2128 let runtime = make_rt().await;
2129 let summary = runtime.import_kg(&archive, &token).await.unwrap();
2130 assert!(summary.skipped_edges.is_empty());
2131 assert_eq!(summary.edges_skipped, 0);
2132 assert_eq!(
2133 serde_json::to_value(&summary).unwrap(),
2134 serde_json::json!({
2135 "entities_imported":entities,"edges_imported":edges,"edges_skipped":0,
2136 "embedding_truncation":{"truncated":0,"discarded_bytes":0}
2137 })
2138 );
2139 }
2140 }
2141
2142 #[test]
2143 fn historical_summary_counts_do_not_invent_skipped_edge_details() {
2144 for count in [0, 2] {
2145 let old = serde_json::json!({
2146 "entities_imported":1,"edges_imported":0,"edges_skipped":count
2147 });
2148 let summary: ImportSummary = serde_json::from_value(old).unwrap();
2149 assert_eq!(summary.edges_skipped, count);
2150 assert!(summary.skipped_edges.is_empty());
2151 assert_eq!(
2152 serde_json::to_value(summary).unwrap(),
2153 serde_json::json!({
2154 "entities_imported":1,"edges_imported":0,"edges_skipped":count,
2155 "embedding_truncation":{"truncated":0,"discarded_bytes":0}
2156 })
2157 );
2158 }
2159 }
2160
2161 #[test]
2162 fn skipped_edge_cause_spellings_are_closed_and_error_is_preserved_verbatim() {
2163 let wire = serde_json::json!({
2164 "edge_id":"00000000-0000-0000-0000-000000000103",
2165 "source":"00000000-0000-0000-0000-000000000001",
2166 "target":"00000000-0000-0000-0000-000000000002",
2167 "relation":"annotates","cause":"endpoint_contract",
2168 "error":"not found: link target changed during endpoint validation"
2169 });
2170 let detail: ImportSkippedEdge = serde_json::from_value(wire.clone()).unwrap();
2171 assert_eq!(detail.cause, ImportEdgeSkipCause::EndpointContract);
2172 assert_eq!(
2173 detail.error.as_deref(),
2174 Some("not found: link target changed during endpoint validation")
2175 );
2176 assert_eq!(serde_json::to_value(detail).unwrap(), wire);
2177 let mut invalid = wire;
2178 invalid["cause"] = serde_json::json!("storage_failure");
2179 assert!(serde_json::from_value::<ImportSkippedEdge>(invalid).is_err());
2180 }
2181}