graphforge_storage/
embedding_discovery.rs

1//! Bounded discovery of compatibility-addressed embedding generations.
2
3use std::path::{Path, PathBuf};
4
5use crate::{
6    EmbeddingCompatibilityDescriptor, EmbeddingCompatibilityId, EmbeddingGenerationPublication,
7    SearchArtifactError, VectorStoreLimits, current_embedding_generation,
8};
9
10/// Maximum compatibility lineages returned by one default discovery.
11pub const MAX_DISCOVERED_EMBEDDING_SPACES: usize = 1_024;
12/// Maximum filesystem entries inspected by one default discovery.
13pub const MAX_EMBEDDING_SPACE_DIRECTORY_ENTRIES: usize = 2_048;
14/// Maximum bytes accepted from one compatibility descriptor.
15pub const MAX_DISCOVERED_EMBEDDING_DESCRIPTOR_BYTES: u64 = 64 * 1024;
16
17const EMBEDDINGS_DIR: &str = "embeddings";
18const SPACES_DIR: &str = "spaces";
19const DESCRIPTOR_FILE: &str = "space.json";
20
21/// Resource bounds for complete embedding-space discovery.
22#[derive(Clone, Copy, Debug)]
23pub struct EmbeddingSpaceDiscoveryLimits {
24    /// Maximum compatibility lineages returned.
25    pub spaces: usize,
26    /// Maximum raw entries inspected under `embeddings/spaces`.
27    pub directory_entries: usize,
28    /// Maximum bytes accepted from one `space.json` descriptor.
29    pub descriptor_bytes: u64,
30    /// Bounds used to validate every active primary vector generation.
31    pub vectors: VectorStoreLimits,
32}
33
34impl Default for EmbeddingSpaceDiscoveryLimits {
35    fn default() -> Self {
36        Self {
37            spaces: MAX_DISCOVERED_EMBEDDING_SPACES,
38            directory_entries: MAX_EMBEDDING_SPACE_DIRECTORY_ENTRIES,
39            descriptor_bytes: MAX_DISCOVERED_EMBEDDING_DESCRIPTOR_BYTES,
40            vectors: VectorStoreLimits::default(),
41        }
42    }
43}
44
45/// One validated compatibility lineage and its optional complete active generation.
46#[derive(Clone, Debug, PartialEq)]
47pub struct DiscoveredEmbeddingSpace {
48    compatibility_id: EmbeddingCompatibilityId,
49    descriptor: EmbeddingCompatibilityDescriptor,
50    active: Option<EmbeddingGenerationPublication>,
51}
52
53impl DiscoveredEmbeddingSpace {
54    /// Exact digest-addressed compatibility identity.
55    #[must_use]
56    pub const fn compatibility_id(&self) -> EmbeddingCompatibilityId {
57        self.compatibility_id
58    }
59
60    /// Fully validated compatibility descriptor reopened from `space.json`.
61    #[must_use]
62    pub const fn descriptor(&self) -> &EmbeddingCompatibilityDescriptor {
63        &self.descriptor
64    }
65
66    /// Fully verified active generation, or `None` before first publication.
67    #[must_use]
68    pub const fn active(&self) -> Option<&EmbeddingGenerationPublication> {
69        self.active.as_ref()
70    }
71}
72
73/// Discover every compatibility-addressed embedding lineage deterministically.
74///
75/// A missing `embeddings/spaces` tree returns an empty vector without creating
76/// directories. Every present lineage is validated from its path through its
77/// descriptor and, when active, through the complete primary-vector reopen path.
78///
79/// # Errors
80/// Returns structured cancellation, limit, corruption, incompatibility, or I/O errors.
81pub fn discover_embedding_spaces<C>(
82    project_dir: &Path,
83    limits: EmbeddingSpaceDiscoveryLimits,
84    mut checkpoint: C,
85) -> Result<Vec<DiscoveredEmbeddingSpace>, SearchArtifactError>
86where
87    C: FnMut() -> Result<(), SearchArtifactError>,
88{
89    validate_limits(limits)?;
90    checkpoint()?;
91    let embeddings = project_dir.join(EMBEDDINGS_DIR);
92    if !path_exists(&embeddings)? {
93        return Ok(Vec::new());
94    }
95    ensure_owned_directory(&embeddings)?;
96    let spaces_root = embeddings.join(SPACES_DIR);
97    if !path_exists(&spaces_root)? {
98        return Ok(Vec::new());
99    }
100    ensure_owned_directory(&spaces_root)?;
101
102    let mut discovered = Vec::new();
103    let mut inspected = 0_usize;
104    let entries = std::fs::read_dir(&spaces_root)
105        .map_err(|source| io("enumerate embedding spaces", &spaces_root, source))?;
106    for entry in entries {
107        checkpoint()?;
108        inspected = inspected.checked_add(1).ok_or_else(|| {
109            exhausted(
110                "embedding_space_directory_entries",
111                limits.directory_entries,
112            )
113        })?;
114        if inspected > limits.directory_entries {
115            return Err(exhausted(
116                "embedding_space_directory_entries",
117                limits.directory_entries,
118            ));
119        }
120        if discovered.len() >= limits.spaces {
121            return Err(exhausted("embedding_spaces", limits.spaces));
122        }
123        let entry =
124            entry.map_err(|source| io("read embedding space entry", &spaces_root, source))?;
125        let path = entry.path();
126        ensure_owned_directory(&path)?;
127        let path_identity = entry
128            .file_name()
129            .to_str()
130            .ok_or_else(|| corrupt(&path, "space directory name is not UTF-8"))
131            .and_then(|value| {
132                EmbeddingCompatibilityId::from_hex(value)
133                    .map_err(|error| corrupt(&path, error.to_string()))
134            })?;
135        if path_exists(&crate::embedding_publication::deletion_marker(
136            project_dir,
137            path_identity,
138        ))? {
139            continue;
140        }
141        let descriptor = read_descriptor(&path, limits.descriptor_bytes)?;
142        let descriptor_identity = descriptor.compatibility_id()?;
143        if descriptor_identity != path_identity {
144            return Err(corrupt(
145                &path,
146                "descriptor compatibility identity does not match its directory",
147            ));
148        }
149        checkpoint()?;
150        let active = current_embedding_generation(
151            project_dir,
152            &descriptor,
153            limits.vectors,
154            &mut checkpoint,
155        )?;
156        discovered.push(DiscoveredEmbeddingSpace {
157            compatibility_id: path_identity,
158            descriptor,
159            active,
160        });
161    }
162    discovered.sort_unstable_by_key(DiscoveredEmbeddingSpace::compatibility_id);
163    Ok(discovered)
164}
165
166fn read_descriptor(
167    space_root: &Path,
168    max_bytes: u64,
169) -> Result<EmbeddingCompatibilityDescriptor, SearchArtifactError> {
170    let path = space_root.join(DESCRIPTOR_FILE);
171    ensure_regular_file(&path)?;
172    let metadata = std::fs::metadata(&path)
173        .map_err(|source| io("inspect embedding descriptor", &path, source))?;
174    if metadata.len() > max_bytes {
175        return Err(SearchArtifactError::ResourceExhausted {
176            resource: "embedding_descriptor_bytes",
177            limit: max_bytes,
178        });
179    }
180    let bytes =
181        std::fs::read(&path).map_err(|source| io("read embedding descriptor", &path, source))?;
182    EmbeddingCompatibilityDescriptor::from_json(&path, &bytes)
183        .map_err(|error| primary_from(space_root, error))
184}
185
186fn validate_limits(limits: EmbeddingSpaceDiscoveryLimits) -> Result<(), SearchArtifactError> {
187    if limits.spaces == 0 || limits.directory_entries == 0 || limits.descriptor_bytes == 0 {
188        Err(invalid(
189            "embedding discovery limits",
190            "must all be non-zero",
191        ))
192    } else {
193        Ok(())
194    }
195}
196
197fn ensure_owned_directory(path: &Path) -> Result<(), SearchArtifactError> {
198    let metadata = std::fs::symlink_metadata(path)
199        .map_err(|source| io("inspect embedding directory", path, source))?;
200    if metadata.file_type().is_symlink() || !metadata.is_dir() {
201        return Err(corrupt(path, "expected an owned directory"));
202    }
203    Ok(())
204}
205
206fn ensure_regular_file(path: &Path) -> Result<(), SearchArtifactError> {
207    let metadata = std::fs::symlink_metadata(path)
208        .map_err(|source| io("inspect embedding descriptor", path, source))?;
209    if metadata.file_type().is_symlink() || !metadata.is_file() {
210        return Err(corrupt(path, "expected a regular descriptor file"));
211    }
212    Ok(())
213}
214
215fn path_exists(path: &Path) -> Result<bool, SearchArtifactError> {
216    match std::fs::symlink_metadata(path) {
217        Ok(_) => Ok(true),
218        Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(false),
219        Err(source) => Err(io("inspect embedding path", path, source)),
220    }
221}
222
223fn primary_from(path: &Path, error: SearchArtifactError) -> SearchArtifactError {
224    match error {
225        SearchArtifactError::IncompatibleManifest { .. }
226        | SearchArtifactError::ResourceExhausted { .. }
227        | SearchArtifactError::Cancelled => error,
228        other => SearchArtifactError::CorruptPrimaryVectors {
229            path: path.to_path_buf(),
230            reason: other.to_string(),
231        },
232    }
233}
234
235fn invalid(field: &'static str, reason: impl Into<String>) -> SearchArtifactError {
236    SearchArtifactError::InvalidSelector {
237        field,
238        reason: reason.into(),
239    }
240}
241
242fn corrupt(path: &Path, reason: impl Into<String>) -> SearchArtifactError {
243    SearchArtifactError::CorruptPrimaryVectors {
244        path: path.to_path_buf(),
245        reason: reason.into(),
246    }
247}
248
249fn exhausted(resource: &'static str, limit: usize) -> SearchArtifactError {
250    SearchArtifactError::ResourceExhausted {
251        resource,
252        limit: limit as u64,
253    }
254}
255
256fn io(operation: &'static str, path: &Path, source: std::io::Error) -> SearchArtifactError {
257    SearchArtifactError::Io {
258        operation,
259        path: PathBuf::from(path),
260        source,
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use std::collections::{BTreeMap, BTreeSet};
267
268    use super::*;
269    use crate::{
270        EmbeddingBatchRow, EmbeddingCompatibilityInput, EmbeddingDistance, EmbeddingNormalization,
271        EmbeddingProducerIdentity, EmbeddingPublicationRequest, EmbeddingSourceState,
272        EmbeddingValueType, SearchCoordinationLimits, VECTOR_DATA_FILE, ValidatedEmbeddingBatch,
273        publish_embedding_generation, validate_embedding_batch,
274    };
275
276    fn descriptor(model: &str) -> EmbeddingCompatibilityDescriptor {
277        EmbeddingCompatibilityDescriptor::new(EmbeddingCompatibilityInput {
278            producer: EmbeddingProducerIdentity::Local {
279                implementation: "discovery-test".to_owned(),
280                model: model.to_owned(),
281                revision: "r1".to_owned(),
282                contract_version: "v1".to_owned(),
283            },
284            dimensions: 2,
285            value_type: EmbeddingValueType::Float32,
286            normalization: EmbeddingNormalization::None,
287            distance: EmbeddingDistance::Cosine,
288            tokenizer: None,
289            chunking: None,
290            hyperparameters: BTreeMap::new(),
291            input_recipe: BTreeMap::from([("property".to_owned(), "body".into())]),
292            source_projection_recipe: BTreeMap::from([("label".to_owned(), "Document".into())]),
293        })
294        .unwrap()
295    }
296
297    fn batch(uuid: [u8; 16]) -> ValidatedEmbeddingBatch {
298        validate_embedding_batch(
299            vec![EmbeddingBatchRow {
300                node_uuid: uuid,
301                vector: vec![1.0, 2.0],
302            }],
303            &BTreeSet::from([uuid]),
304            2,
305            EmbeddingNormalization::None,
306            VectorStoreLimits::default(),
307            || Ok(()),
308        )
309        .unwrap()
310    }
311
312    fn publish(
313        project: &Path,
314        descriptor: &EmbeddingCompatibilityDescriptor,
315        marker: u8,
316    ) -> PathBuf {
317        publish_embedding_generation(
318            project,
319            EmbeddingPublicationRequest {
320                descriptor,
321                source: EmbeddingSourceState::new(1, [marker; 32], [marker + 1; 32], 1),
322                batch: &batch([marker; 16]),
323                generated_at_micros: 10,
324                committed_at_micros: 11,
325            },
326            VectorStoreLimits::default(),
327            SearchCoordinationLimits::default(),
328            || Ok(()),
329        )
330        .unwrap()
331        .publication()
332        .path
333        .clone()
334    }
335
336    fn descriptor_only(project: &Path, descriptor: &EmbeddingCompatibilityDescriptor) {
337        let identity = descriptor.compatibility_id().unwrap();
338        let root = project
339            .join(EMBEDDINGS_DIR)
340            .join(SPACES_DIR)
341            .join(identity.to_hex());
342        std::fs::create_dir_all(&root).unwrap();
343        std::fs::write(
344            root.join(DESCRIPTOR_FILE),
345            descriptor.to_canonical_json().unwrap(),
346        )
347        .unwrap();
348    }
349
350    #[test]
351    fn missing_and_empty_trees_return_no_spaces_without_creating_files() {
352        let project = tempfile::tempdir().unwrap();
353        assert!(
354            discover_embedding_spaces(
355                project.path(),
356                EmbeddingSpaceDiscoveryLimits::default(),
357                || Ok(())
358            )
359            .unwrap()
360            .is_empty()
361        );
362        assert!(!project.path().join(EMBEDDINGS_DIR).exists());
363        std::fs::create_dir_all(project.path().join(EMBEDDINGS_DIR).join(SPACES_DIR)).unwrap();
364        assert!(
365            discover_embedding_spaces(
366                project.path(),
367                EmbeddingSpaceDiscoveryLimits::default(),
368                || Ok(())
369            )
370            .unwrap()
371            .is_empty()
372        );
373    }
374
375    #[test]
376    fn active_and_descriptor_only_lineages_are_sorted_and_reopen_completely() {
377        let project = tempfile::tempdir().unwrap();
378        let a = descriptor("a");
379        let b = descriptor("b");
380        let _ = publish(project.path(), &b, 2);
381        descriptor_only(project.path(), &a);
382
383        let discovered = discover_embedding_spaces(
384            project.path(),
385            EmbeddingSpaceDiscoveryLimits::default(),
386            || Ok(()),
387        )
388        .unwrap();
389        assert_eq!(discovered.len(), 2);
390        assert!(
391            discovered
392                .windows(2)
393                .all(|pair| pair[0].compatibility_id() < pair[1].compatibility_id())
394        );
395        let active = discovered
396            .iter()
397            .find(|space| space.descriptor() == &b)
398            .unwrap()
399            .active()
400            .unwrap();
401        assert_eq!(active.manifest.vector_count(), 1);
402        assert!(
403            discovered
404                .iter()
405                .find(|space| space.descriptor() == &a)
406                .unwrap()
407                .active()
408                .is_none()
409        );
410    }
411
412    #[test]
413    fn hostile_entries_identity_corruption_limits_and_cancellation_fail_closed() {
414        let project = tempfile::tempdir().unwrap();
415        let stable = descriptor("stable");
416        descriptor_only(project.path(), &stable);
417        let spaces = project.path().join(EMBEDDINGS_DIR).join(SPACES_DIR);
418
419        std::fs::write(spaces.join("not-a-space"), b"hostile").unwrap();
420        assert!(matches!(
421            discover_embedding_spaces(
422                project.path(),
423                EmbeddingSpaceDiscoveryLimits::default(),
424                || Ok(())
425            ),
426            Err(SearchArtifactError::CorruptPrimaryVectors { .. })
427        ));
428        std::fs::remove_file(spaces.join("not-a-space")).unwrap();
429
430        let wrong = descriptor("wrong");
431        let stable_root = spaces.join(stable.compatibility_id().unwrap().to_hex());
432        std::fs::write(
433            stable_root.join(DESCRIPTOR_FILE),
434            wrong.to_canonical_json().unwrap(),
435        )
436        .unwrap();
437        assert!(matches!(
438            discover_embedding_spaces(
439                project.path(),
440                EmbeddingSpaceDiscoveryLimits::default(),
441                || Ok(())
442            ),
443            Err(SearchArtifactError::CorruptPrimaryVectors { .. })
444        ));
445        std::fs::write(
446            stable_root.join(DESCRIPTOR_FILE),
447            stable.to_canonical_json().unwrap(),
448        )
449        .unwrap();
450
451        assert!(matches!(
452            discover_embedding_spaces(
453                project.path(),
454                EmbeddingSpaceDiscoveryLimits {
455                    spaces: 1,
456                    directory_entries: 1,
457                    descriptor_bytes: 8,
458                    vectors: VectorStoreLimits::default(),
459                },
460                || Ok(())
461            ),
462            Err(SearchArtifactError::ResourceExhausted { .. })
463        ));
464        assert!(matches!(
465            discover_embedding_spaces(
466                project.path(),
467                EmbeddingSpaceDiscoveryLimits::default(),
468                || Err(SearchArtifactError::Cancelled)
469            ),
470            Err(SearchArtifactError::Cancelled)
471        ));
472    }
473
474    #[test]
475    fn corrupt_active_primary_vectors_are_never_reported_as_absent() {
476        let project = tempfile::tempdir().unwrap();
477        let descriptor = descriptor("corrupt-primary");
478        let generation = publish(project.path(), &descriptor, 4);
479        std::fs::write(generation.join(VECTOR_DATA_FILE), b"corrupt").unwrap();
480        assert!(matches!(
481            discover_embedding_spaces(
482                project.path(),
483                EmbeddingSpaceDiscoveryLimits::default(),
484                || Ok(())
485            ),
486            Err(SearchArtifactError::CorruptPrimaryVectors { .. })
487        ));
488    }
489
490    #[cfg(unix)]
491    #[test]
492    fn symlinked_lineage_fails_closed() {
493        use std::os::unix::fs::symlink;
494
495        let project = tempfile::tempdir().unwrap();
496        let spaces = project.path().join(EMBEDDINGS_DIR).join(SPACES_DIR);
497        std::fs::create_dir_all(&spaces).unwrap();
498        let target = tempfile::tempdir().unwrap();
499        symlink(target.path(), spaces.join("0".repeat(64))).unwrap();
500        assert!(matches!(
501            discover_embedding_spaces(
502                project.path(),
503                EmbeddingSpaceDiscoveryLimits::default(),
504                || Ok(())
505            ),
506            Err(SearchArtifactError::CorruptPrimaryVectors { .. })
507        ));
508    }
509}
graphforge_storage/embedding_discovery.rs

graphforge_storage/
embedding_discovery.rs
