dag_ml_data_core/
fitted_adapter.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::fs;
3use std::path::Path;
4use std::sync::Mutex;
5
6use serde::{Deserialize, Serialize};
7
8use crate::error::{DataError, Result};
9
10pub const FITTED_ADAPTER_REF_SCHEMA_VERSION: u32 = 1;
11pub const FITTED_ADAPTER_REF_SCHEMA_ID: &str =
12    "https://github.com/GBeurier/dag-ml-data/schemas/fitted_adapter_ref.v1.schema.json";
13pub const FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION: u32 = 1;
14
15#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
16#[serde(rename_all = "snake_case")]
17pub enum FittedAdapterBackend {
18    Joblib,
19    Pickle,
20    Json,
21    Numpy,
22    Onnx,
23    Raw,
24}
25
26#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
27pub struct FittedAdapterRef {
28    #[serde(default = "default_fitted_adapter_ref_schema_version")]
29    pub schema_version: u32,
30    pub adapter_id: String,
31    pub adapter_version: String,
32    pub params_fingerprint: String,
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub backend: Option<FittedAdapterBackend>,
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub uri: Option<String>,
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub content_fingerprint: Option<String>,
39    // Matches `dag-ml` `ArtifactRef.size_bytes` wire shape: when `None` this
40    // field serializes as `"size_bytes": null`, not omitted, so the two refs
41    // round-trip identically over JSON.
42    #[serde(default)]
43    pub size_bytes: Option<u64>,
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub plugin: Option<String>,
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub plugin_version: Option<String>,
48    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
49    pub metadata: BTreeMap<String, serde_json::Value>,
50}
51
52fn default_fitted_adapter_ref_schema_version() -> u32 {
53    FITTED_ADAPTER_REF_SCHEMA_VERSION
54}
55
56fn default_fitted_adapter_manifest_schema_version() -> u32 {
57    FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION
58}
59
60impl FittedAdapterRef {
61    pub fn validate(&self) -> Result<()> {
62        if self.schema_version != FITTED_ADAPTER_REF_SCHEMA_VERSION {
63            return Err(DataError::Validation(format!(
64                "fitted adapter `{}` uses unsupported schema_version {}, expected {}",
65                self.adapter_id, self.schema_version, FITTED_ADAPTER_REF_SCHEMA_VERSION
66            )));
67        }
68        if self.adapter_id.trim().is_empty() {
69            return Err(DataError::Validation(
70                "fitted adapter ref has empty adapter_id".to_string(),
71            ));
72        }
73        if self.adapter_version.trim().is_empty() {
74            return Err(DataError::Validation(format!(
75                "fitted adapter `{}` has empty adapter_version",
76                self.adapter_id
77            )));
78        }
79        validate_fingerprint(
80            "fitted adapter params",
81            &self.params_fingerprint,
82            &self.adapter_id,
83        )?;
84        validate_optional_text("uri", &self.uri, &self.adapter_id)?;
85        validate_optional_text("plugin", &self.plugin, &self.adapter_id)?;
86        validate_optional_text("plugin_version", &self.plugin_version, &self.adapter_id)?;
87        if self.plugin_version.is_some() && self.plugin.is_none() {
88            return Err(DataError::Validation(format!(
89                "fitted adapter `{}` has plugin_version without plugin",
90                self.adapter_id
91            )));
92        }
93        if let Some(content_fingerprint) = &self.content_fingerprint {
94            validate_fingerprint(
95                "fitted adapter content",
96                content_fingerprint,
97                &self.adapter_id,
98            )?;
99        }
100        if self.uri.is_some() && self.backend.is_none() {
101            return Err(DataError::Validation(format!(
102                "fitted adapter `{}` has uri without backend",
103                self.adapter_id
104            )));
105        }
106        if self.uri.is_some() && self.content_fingerprint.is_none() {
107            return Err(DataError::Validation(format!(
108                "fitted adapter `{}` has uri without content_fingerprint",
109                self.adapter_id
110            )));
111        }
112        for key in self.metadata.keys() {
113            if key.trim().is_empty() {
114                return Err(DataError::Validation(format!(
115                    "fitted adapter `{}` metadata contains an empty key",
116                    self.adapter_id
117                )));
118            }
119        }
120        Ok(())
121    }
122
123    /// Validate that the fitted adapter carries portable persistence metadata:
124    /// a backend, a safe relative URI and a content fingerprint. Legacy refs
125    /// that only carry inline state stay readable through [`Self::validate`]
126    /// but are refused here so persisted manifests can be moved with their
127    /// payloads.
128    pub fn validate_portable(&self) -> Result<()> {
129        self.validate()?;
130        let Some(uri) = self.uri.as_deref() else {
131            return Err(DataError::Validation(format!(
132                "fitted adapter `{}` is not portable: requires backend, uri and content_fingerprint",
133                self.adapter_id
134            )));
135        };
136        validate_relative_uri(&self.adapter_id, uri)
137    }
138}
139
140#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
141pub struct FittedAdapterManifestEntry {
142    pub adapter_id: String,
143    pub fitted_adapter: FittedAdapterRef,
144}
145
146#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
147pub struct FittedAdapterManifest {
148    #[serde(default = "default_fitted_adapter_manifest_schema_version")]
149    pub schema_version: u32,
150    pub entries: Vec<FittedAdapterManifestEntry>,
151}
152
153impl FittedAdapterManifest {
154    pub fn new(entries: Vec<FittedAdapterManifestEntry>) -> Self {
155        Self {
156            schema_version: FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION,
157            entries,
158        }
159    }
160
161    pub fn validate(&self) -> Result<()> {
162        self.validate_inner(false)
163    }
164
165    /// Validate the manifest and require every entry to carry portable
166    /// persistence metadata. Use this when serializing manifests for replay
167    /// across machines or repositories.
168    pub fn validate_portable(&self) -> Result<()> {
169        self.validate_inner(true)
170    }
171
172    /// Writes the manifest as pretty-printed UTF-8 JSON to `path`.
173    /// `require_portable` toggles between [`Self::validate`] and
174    /// [`Self::validate_portable`] before writing, so a manifest is never
175    /// persisted in a state that would not load cleanly on the other side.
176    pub fn write_to_path(&self, path: &Path, require_portable: bool) -> Result<()> {
177        if require_portable {
178            self.validate_portable()?;
179        } else {
180            self.validate()?;
181        }
182        let payload = serde_json::to_vec_pretty(self).map_err(|error| {
183            DataError::Validation(format!(
184                "failed to serialize fitted adapter manifest to JSON: {error}"
185            ))
186        })?;
187        fs::write(path, payload).map_err(|error| {
188            DataError::Validation(format!(
189                "failed to write fitted adapter manifest to `{}`: {error}",
190                path.display()
191            ))
192        })
193    }
194
195    /// Reads the manifest back from `path`. Applies the same `require_portable`
196    /// gate as [`Self::write_to_path`] so loaded manifests are immediately
197    /// safe to use.
198    pub fn read_from_path(path: &Path, require_portable: bool) -> Result<Self> {
199        let bytes = fs::read(path).map_err(|error| {
200            DataError::Validation(format!(
201                "failed to read fitted adapter manifest from `{}`: {error}",
202                path.display()
203            ))
204        })?;
205        let manifest: Self = serde_json::from_slice(&bytes).map_err(|error| {
206            DataError::Validation(format!(
207                "failed to parse fitted adapter manifest from `{}`: {error}",
208                path.display()
209            ))
210        })?;
211        if require_portable {
212            manifest.validate_portable()?;
213        } else {
214            manifest.validate()?;
215        }
216        Ok(manifest)
217    }
218
219    fn validate_inner(&self, require_portable: bool) -> Result<()> {
220        if self.schema_version != FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION {
221            return Err(DataError::Validation(format!(
222                "fitted adapter manifest uses unsupported schema_version {}, expected {}",
223                self.schema_version, FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION
224            )));
225        }
226        let mut seen = BTreeSet::new();
227        for entry in &self.entries {
228            if entry.adapter_id != entry.fitted_adapter.adapter_id {
229                return Err(DataError::Validation(format!(
230                    "fitted adapter manifest entry key `{}` does not match ref adapter_id `{}`",
231                    entry.adapter_id, entry.fitted_adapter.adapter_id
232                )));
233            }
234            if !seen.insert(&entry.adapter_id) {
235                return Err(DataError::Validation(format!(
236                    "fitted adapter manifest contains duplicate adapter_id `{}`",
237                    entry.adapter_id
238                )));
239            }
240            if require_portable {
241                entry.fitted_adapter.validate_portable()?;
242            } else {
243                entry.fitted_adapter.validate()?;
244            }
245        }
246        Ok(())
247    }
248}
249
250#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
251pub struct FittedAdapterMaterializationRequest {
252    pub adapter_id: String,
253    pub params_fingerprint: String,
254}
255
256impl FittedAdapterMaterializationRequest {
257    pub fn validate(&self) -> Result<()> {
258        if self.adapter_id.trim().is_empty() {
259            return Err(DataError::Validation(
260                "fitted adapter materialization request has empty adapter_id".to_string(),
261            ));
262        }
263        validate_fingerprint(
264            "fitted adapter materialization params",
265            &self.params_fingerprint,
266            &self.adapter_id,
267        )?;
268        Ok(())
269    }
270}
271
272#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
273pub struct FittedAdapterHandleRecord {
274    pub handle: u64,
275    pub fitted_adapter: FittedAdapterRef,
276}
277
278/// Runtime contract for materializing opaque host-owned fitted-adapter
279/// handles from a previously registered `FittedAdapterRef`. The store
280/// validates the request's `params_fingerprint` against the registered ref
281/// and returns the opaque handle id. It never reads, writes or deserializes
282/// adapter payloads — payload materialization happens host-side, behind the
283/// opaque handle.
284pub trait RuntimeFittedAdapterStore {
285    fn materialize(&self, request: &FittedAdapterMaterializationRequest) -> Result<u64>;
286}
287
288#[derive(Debug, Default)]
289pub struct InMemoryFittedAdapterStore {
290    state: Mutex<FittedAdapterStoreState>,
291}
292
293#[derive(Debug, Default)]
294struct FittedAdapterStoreState {
295    next_handle: u64,
296    records: BTreeMap<String, FittedAdapterHandleRecord>,
297}
298
299impl InMemoryFittedAdapterStore {
300    pub fn new() -> Self {
301        Self {
302            state: Mutex::new(FittedAdapterStoreState {
303                next_handle: 1,
304                records: BTreeMap::new(),
305            }),
306        }
307    }
308
309    /// Locks the store, propagating poison errors as a validation failure.
310    /// Poisoning happens when another thread panicked while holding the lock;
311    /// for an in-memory contract store the right policy is to surface the
312    /// error rather than silently `recover` from it.
313    fn lock(&self) -> Result<std::sync::MutexGuard<'_, FittedAdapterStoreState>> {
314        self.state.lock().map_err(|error| {
315            DataError::Validation(format!("fitted adapter store mutex poisoned: {error}"))
316        })
317    }
318
319    pub fn register(&self, fitted_adapter: FittedAdapterRef) -> Result<FittedAdapterHandleRecord> {
320        fitted_adapter.validate()?;
321        let mut state = self.lock()?;
322        if state.records.contains_key(&fitted_adapter.adapter_id) {
323            return Err(DataError::Validation(format!(
324                "fitted adapter store already has handle for `{}`",
325                fitted_adapter.adapter_id
326            )));
327        }
328        let handle = state.next_handle;
329        state.next_handle += 1;
330        let record = FittedAdapterHandleRecord {
331            handle,
332            fitted_adapter,
333        };
334        state
335            .records
336            .insert(record.fitted_adapter.adapter_id.clone(), record.clone());
337        Ok(record)
338    }
339
340    pub fn register_manifest(
341        &self,
342        manifest: &FittedAdapterManifest,
343    ) -> Result<Vec<FittedAdapterHandleRecord>> {
344        manifest.validate()?;
345        manifest
346            .entries
347            .iter()
348            .map(|entry| self.register(entry.fitted_adapter.clone()))
349            .collect()
350    }
351
352    pub fn get(&self, adapter_id: &str) -> Option<FittedAdapterHandleRecord> {
353        let state = self.lock().ok()?;
354        state.records.get(adapter_id).cloned()
355    }
356
357    pub fn release(&self, adapter_id: &str) -> bool {
358        match self.lock() {
359            Ok(mut state) => state.records.remove(adapter_id).is_some(),
360            Err(_) => false,
361        }
362    }
363
364    pub fn len(&self) -> usize {
365        self.lock().map(|state| state.records.len()).unwrap_or(0)
366    }
367
368    pub fn is_empty(&self) -> bool {
369        self.lock()
370            .map(|state| state.records.is_empty())
371            .unwrap_or(true)
372    }
373}
374
375impl RuntimeFittedAdapterStore for InMemoryFittedAdapterStore {
376    fn materialize(&self, request: &FittedAdapterMaterializationRequest) -> Result<u64> {
377        request.validate()?;
378        let state = self.lock()?;
379        let record = state.records.get(&request.adapter_id).ok_or_else(|| {
380            DataError::Validation(format!(
381                "fitted adapter store is missing adapter `{}`",
382                request.adapter_id
383            ))
384        })?;
385        if record.fitted_adapter.params_fingerprint != request.params_fingerprint {
386            return Err(DataError::FingerprintMismatch {
387                kind: "params",
388                expected: record.fitted_adapter.params_fingerprint.clone(),
389                actual: request.params_fingerprint.clone(),
390            });
391        }
392        Ok(record.handle)
393    }
394}
395
396fn validate_optional_text(label: &str, value: &Option<String>, adapter_id: &str) -> Result<()> {
397    if let Some(text) = value {
398        if text.trim().is_empty() {
399            return Err(DataError::Validation(format!(
400                "fitted adapter `{adapter_id}` has empty {label}"
401            )));
402        }
403        if text.chars().any(char::is_control) {
404            return Err(DataError::Validation(format!(
405                "fitted adapter `{adapter_id}` has control characters in {label}"
406            )));
407        }
408    }
409    Ok(())
410}
411
412fn validate_fingerprint(label: &str, value: &str, adapter_id: &str) -> Result<()> {
413    if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
414        return Err(DataError::Validation(format!(
415            "fitted adapter `{adapter_id}` {label} fingerprint must be a 64-character hex digest"
416        )));
417    }
418    Ok(())
419}
420
421/// Deterministic path safety for relative fitted-adapter URIs. Rejects empty
422/// values, control characters, absolute paths (POSIX root, Windows root or
423/// drive prefix), URI schemes such as `http://`, `s3://` or `file://` (any
424/// colon in the leading path segment) and any `..` traversal component.
425/// Parsing is platform-independent so portable manifests validate identically
426/// everywhere; it adds no dependency. Kept byte-for-byte equivalent with
427/// `dag-ml`'s `validate_relative_artifact_uri` in
428/// `crates/dag-ml-core/src/runtime/prediction_store.rs` so portable refs accepted by one repo
429/// are accepted by the other.
430fn validate_relative_uri(adapter_id: &str, uri: &str) -> Result<()> {
431    if uri.is_empty() {
432        return Err(DataError::Validation(format!(
433            "fitted adapter `{adapter_id}` has empty uri"
434        )));
435    }
436    if uri.chars().any(char::is_control) {
437        return Err(DataError::Validation(format!(
438            "fitted adapter `{adapter_id}` uri has control characters"
439        )));
440    }
441    if uri.starts_with('/') || uri.starts_with('\\') {
442        return Err(DataError::Validation(format!(
443            "fitted adapter `{adapter_id}` uri `{uri}` must be a relative path"
444        )));
445    }
446    let mut prefix = uri.chars();
447    if let (Some(drive), Some(':')) = (prefix.next(), prefix.next()) {
448        if drive.is_ascii_alphabetic() {
449            return Err(DataError::Validation(format!(
450                "fitted adapter `{adapter_id}` uri `{uri}` must be a relative path"
451            )));
452        }
453    }
454    let first_segment = uri.split(['/', '\\']).next().unwrap_or(uri);
455    if first_segment.contains(':') {
456        return Err(DataError::Validation(format!(
457            "fitted adapter `{adapter_id}` uri `{uri}` must not include a scheme or colon in its first path segment"
458        )));
459    }
460    for segment in uri.split(['/', '\\']) {
461        if segment == ".." {
462            return Err(DataError::Validation(format!(
463                "fitted adapter `{adapter_id}` uri `{uri}` must not contain `..` components"
464            )));
465        }
466    }
467    Ok(())
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473
474    fn fingerprint(byte: u8) -> String {
475        format!("{byte:02x}").repeat(32)
476    }
477
478    fn portable_ref() -> FittedAdapterRef {
479        FittedAdapterRef {
480            schema_version: FITTED_ADAPTER_REF_SCHEMA_VERSION,
481            adapter_id: "snv".to_string(),
482            adapter_version: "1.0.0".to_string(),
483            params_fingerprint: fingerprint(0x12),
484            backend: Some(FittedAdapterBackend::Joblib),
485            uri: Some("fitted/snv.joblib".to_string()),
486            content_fingerprint: Some(fingerprint(0x34)),
487            size_bytes: Some(2048),
488            plugin: Some("sklearn".to_string()),
489            plugin_version: Some("1.4.0".to_string()),
490            metadata: BTreeMap::new(),
491        }
492    }
493
494    #[test]
495    fn portable_ref_validates() {
496        let value = portable_ref();
497        value.validate().unwrap();
498        value.validate_portable().unwrap();
499    }
500
501    #[test]
502    fn inline_ref_validates_but_is_not_portable() {
503        let value = FittedAdapterRef {
504            uri: None,
505            backend: None,
506            content_fingerprint: None,
507            size_bytes: None,
508            plugin: None,
509            plugin_version: None,
510            ..portable_ref()
511        };
512        value.validate().unwrap();
513        let error = value.validate_portable().unwrap_err();
514        assert!(format!("{error}").contains("is not portable"));
515    }
516
517    #[test]
518    fn rejects_unsupported_schema_version() {
519        let mut value = portable_ref();
520        value.schema_version = FITTED_ADAPTER_REF_SCHEMA_VERSION + 1;
521        let error = value.validate().unwrap_err();
522        assert!(format!("{error}").contains("unsupported schema_version"));
523    }
524
525    #[test]
526    fn rejects_uri_without_backend_or_fingerprint() {
527        let mut value = portable_ref();
528        value.backend = None;
529        let error = value.validate().unwrap_err();
530        assert!(format!("{error}").contains("has uri without backend"));
531
532        let mut value = portable_ref();
533        value.content_fingerprint = None;
534        let error = value.validate().unwrap_err();
535        assert!(format!("{error}").contains("has uri without content_fingerprint"));
536    }
537
538    #[test]
539    fn rejects_plugin_version_without_plugin() {
540        let mut value = portable_ref();
541        value.plugin = None;
542        let error = value.validate().unwrap_err();
543        assert!(format!("{error}").contains("has plugin_version without plugin"));
544    }
545
546    /// Locks down the exact accept/reject behaviour of the portable URI
547    /// rules so the byte-for-byte copy in dag-ml's
548    /// `validate_relative_artifact_uri` can be cross-checked. Any change to
549    /// either function must also update this fixture; the script-level check
550    /// in `scripts/validate_contracts.py` reads both function bodies and
551    /// asserts textual equivalence as a second line of defense.
552    #[test]
553    fn portable_uri_rules_are_locked_to_a_shared_fixture() {
554        let accept = [
555            "snv.joblib",
556            "fitted/snv.joblib",
557            "nested/dir.name/file-v1.bin",
558            "alpha_beta/gamma.json",
559            "a.b/c.d",
560        ];
561        let reject_absolute = ["/abs/path.joblib", "\\abs\\path.joblib"];
562        let reject_drive = ["C:\\fitted\\snv.joblib", "Z:/path"];
563        let reject_schemes = [
564            "http://example.com/x.bin",
565            "s3://bucket/x.bin",
566            "file:///x.bin",
567            "scheme:relative",
568        ];
569        let reject_traversal = [
570            "../escape.bin",
571            "fitted/../escape.bin",
572            "ok/../../escape.bin",
573        ];
574        let reject_empty_or_control = ["", "ctrl/\u{0001}path", "\n/path"];
575
576        for uri in accept {
577            let value = FittedAdapterRef {
578                uri: Some(uri.to_string()),
579                ..portable_ref()
580            };
581            if let Err(error) = value.validate_portable() {
582                panic!("expected `{uri}` to validate portable: {error}");
583            }
584        }
585        for group in [
586            &reject_absolute[..],
587            &reject_drive[..],
588            &reject_schemes[..],
589            &reject_traversal[..],
590            &reject_empty_or_control[..],
591        ] {
592            for uri in group {
593                let value = FittedAdapterRef {
594                    uri: Some(uri.to_string()),
595                    ..portable_ref()
596                };
597                let error = value
598                    .validate_portable()
599                    .expect_err(&format!("expected `{uri}` to be rejected as non-portable"));
600                let message = format!("{error}");
601                assert!(
602                    !message.is_empty(),
603                    "rejection of `{uri}` must carry a non-empty error message"
604                );
605            }
606        }
607    }
608
609    #[test]
610    fn rejects_non_portable_uris() {
611        for uri in [
612            "/abs/path.joblib",
613            "C:\\fitted\\snv.joblib",
614            "http://example.com/fitted/snv.joblib",
615            "s3://bucket/fitted/snv.joblib",
616            "file:///fitted/snv.joblib",
617            "../escape/snv.joblib",
618            "fitted/../escape/snv.joblib",
619        ] {
620            let value = FittedAdapterRef {
621                uri: Some(uri.to_string()),
622                ..portable_ref()
623            };
624            let error = value.validate_portable().unwrap_err();
625            assert!(
626                format!("{error}").contains("must") || format!("{error}").contains("relative path"),
627                "URI `{uri}` should be rejected, got: {error}"
628            );
629        }
630    }
631
632    #[test]
633    fn rejects_uri_with_control_chars_or_empty() {
634        let mut value = portable_ref();
635        value.uri = Some("".to_string());
636        let error = value.validate_portable().unwrap_err();
637        assert!(format!("{error}").contains("empty uri"));
638
639        let mut value = portable_ref();
640        value.uri = Some("fitted/\u{0001}".to_string());
641        let error = value.validate_portable().unwrap_err();
642        assert!(format!("{error}").contains("control characters"));
643    }
644
645    #[test]
646    fn rejects_whitespace_only_optional_text_fields() {
647        for (label, value) in [
648            (
649                "uri",
650                FittedAdapterRef {
651                    uri: Some("   ".to_string()),
652                    ..portable_ref()
653                },
654            ),
655            (
656                "plugin",
657                FittedAdapterRef {
658                    plugin: Some("\t".to_string()),
659                    ..portable_ref()
660                },
661            ),
662            (
663                "plugin_version",
664                FittedAdapterRef {
665                    plugin_version: Some("   ".to_string()),
666                    ..portable_ref()
667                },
668            ),
669        ] {
670            let error = value.validate().unwrap_err();
671            let message = format!("{error}");
672            assert!(
673                message.contains(&format!("has empty {label}")),
674                "expected empty {label} error, got: {message}"
675            );
676        }
677    }
678
679    #[test]
680    fn published_fitted_adapter_ref_schema_pins_current_id_and_backends() {
681        let schema: serde_json::Value = serde_json::from_str(include_str!(
682            "../../../docs/contracts/fitted_adapter_ref.schema.json"
683        ))
684        .unwrap();
685        assert_eq!(schema["$id"].as_str(), Some(FITTED_ADAPTER_REF_SCHEMA_ID));
686        assert!(
687            schema["$id"]
688                .as_str()
689                .unwrap()
690                .ends_with(&format!("v{FITTED_ADAPTER_REF_SCHEMA_VERSION}.schema.json")),
691            "schema $id `{}` must encode version v{FITTED_ADAPTER_REF_SCHEMA_VERSION}",
692            schema["$id"]
693        );
694        let backends = schema["$defs"]["backend"]["enum"].as_array().unwrap();
695        for expected in ["joblib", "pickle", "json", "numpy", "onnx", "raw"] {
696            assert!(
697                backends
698                    .iter()
699                    .any(|value| value.as_str() == Some(expected)),
700                "fitted_adapter_ref backend enum is missing `{expected}`"
701            );
702        }
703    }
704
705    #[test]
706    fn rejects_metadata_with_empty_keys() {
707        let mut value = portable_ref();
708        value
709            .metadata
710            .insert("".to_string(), serde_json::Value::Null);
711        let error = value.validate().unwrap_err();
712        assert!(format!("{error}").contains("metadata contains an empty key"));
713    }
714
715    #[test]
716    fn rejects_invalid_fingerprints() {
717        let mut value = portable_ref();
718        value.params_fingerprint = "not-a-fingerprint".to_string();
719        let error = value.validate().unwrap_err();
720        assert!(format!("{error}").contains("params fingerprint"));
721
722        let mut value = portable_ref();
723        value.content_fingerprint = Some("not-a-fingerprint".to_string());
724        let error = value.validate().unwrap_err();
725        assert!(format!("{error}").contains("content fingerprint"));
726    }
727
728    #[test]
729    fn manifest_validates_and_requires_unique_adapter_ids() {
730        let entry = FittedAdapterManifestEntry {
731            adapter_id: "snv".to_string(),
732            fitted_adapter: portable_ref(),
733        };
734        let manifest = FittedAdapterManifest::new(vec![entry.clone()]);
735        manifest.validate().unwrap();
736        manifest.validate_portable().unwrap();
737
738        let manifest = FittedAdapterManifest::new(vec![entry.clone(), entry]);
739        let error = manifest.validate().unwrap_err();
740        assert!(format!("{error}").contains("duplicate adapter_id"));
741    }
742
743    #[test]
744    fn manifest_refuses_key_mismatch_and_bad_schema_version() {
745        let entry = FittedAdapterManifestEntry {
746            adapter_id: "snv".to_string(),
747            fitted_adapter: FittedAdapterRef {
748                adapter_id: "msc".to_string(),
749                ..portable_ref()
750            },
751        };
752        let manifest = FittedAdapterManifest::new(vec![entry]);
753        let error = manifest.validate().unwrap_err();
754        assert!(format!("{error}").contains("does not match ref adapter_id"));
755
756        let mut manifest = FittedAdapterManifest::new(vec![]);
757        manifest.schema_version = FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION + 1;
758        let error = manifest.validate().unwrap_err();
759        assert!(format!("{error}").contains("unsupported schema_version"));
760    }
761
762    #[test]
763    fn manifest_round_trips_through_file_persistence() {
764        let manifest = FittedAdapterManifest::new(vec![FittedAdapterManifestEntry {
765            adapter_id: "snv".to_string(),
766            fitted_adapter: portable_ref(),
767        }]);
768        let path = std::env::temp_dir().join(format!(
769            "dag_ml_data_fitted_adapter_manifest_roundtrip_{}.json",
770            std::process::id()
771        ));
772        manifest.write_to_path(&path, true).unwrap();
773        let loaded = FittedAdapterManifest::read_from_path(&path, true).unwrap();
774        assert_eq!(loaded, manifest);
775        let _ = std::fs::remove_file(&path);
776    }
777
778    #[test]
779    fn manifest_file_persistence_rejects_inline_when_portable_required() {
780        let manifest = FittedAdapterManifest::new(vec![FittedAdapterManifestEntry {
781            adapter_id: "snv".to_string(),
782            fitted_adapter: FittedAdapterRef {
783                uri: None,
784                backend: None,
785                content_fingerprint: None,
786                plugin: None,
787                plugin_version: None,
788                ..portable_ref()
789            },
790        }]);
791        let path = std::env::temp_dir().join(format!(
792            "dag_ml_data_fitted_adapter_manifest_inline_{}.json",
793            std::process::id()
794        ));
795        let error = manifest.write_to_path(&path, true).unwrap_err();
796        assert!(format!("{error}").contains("is not portable"));
797        manifest.write_to_path(&path, false).unwrap();
798        let loaded_inline = FittedAdapterManifest::read_from_path(&path, false).unwrap();
799        assert_eq!(loaded_inline, manifest);
800        let portable_load_error = FittedAdapterManifest::read_from_path(&path, true).unwrap_err();
801        assert!(format!("{portable_load_error}").contains("is not portable"));
802        let _ = std::fs::remove_file(&path);
803    }
804
805    #[test]
806    fn manifest_file_persistence_reports_missing_paths() {
807        let path = std::env::temp_dir().join(format!(
808            "dag_ml_data_fitted_adapter_manifest_missing_{}.json",
809            std::process::id()
810        ));
811        // Ensure path does not exist.
812        let _ = std::fs::remove_file(&path);
813        let error = FittedAdapterManifest::read_from_path(&path, false).unwrap_err();
814        assert!(format!("{error}").contains("failed to read fitted adapter manifest"));
815    }
816
817    #[test]
818    fn in_memory_store_is_sync_and_handles_concurrent_register_calls() {
819        use std::sync::Arc;
820
821        // Compile-time check: store must be Send + Sync so it can be shared
822        // across host threads through the C ABI handle.
823        fn assert_send_sync<T: Send + Sync>() {}
824        assert_send_sync::<InMemoryFittedAdapterStore>();
825
826        let store = Arc::new(InMemoryFittedAdapterStore::new());
827        let mut handles = Vec::new();
828        for thread_idx in 0..8 {
829            let store = store.clone();
830            handles.push(std::thread::spawn(move || {
831                let mut record = portable_ref();
832                record.adapter_id = format!("snv:{thread_idx}");
833                store.register(record).unwrap()
834            }));
835        }
836        let records: Vec<_> = handles
837            .into_iter()
838            .map(|handle| handle.join().unwrap())
839            .collect();
840        assert_eq!(records.len(), 8);
841        assert_eq!(store.len(), 8);
842        let mut handle_ids: Vec<u64> = records.iter().map(|record| record.handle).collect();
843        handle_ids.sort();
844        // Handles are assigned 1..=8 across the threads, no duplicates.
845        assert_eq!(handle_ids, (1..=8).collect::<Vec<u64>>());
846    }
847
848    #[test]
849    fn in_memory_store_registers_and_materializes_adapter_handles() {
850        let store = InMemoryFittedAdapterStore::new();
851        assert!(store.is_empty());
852
853        let record = store.register(portable_ref()).unwrap();
854        assert_eq!(record.handle, 1);
855        assert_eq!(record.fitted_adapter.adapter_id, "snv");
856        assert_eq!(store.len(), 1);
857
858        let request = FittedAdapterMaterializationRequest {
859            adapter_id: "snv".to_string(),
860            params_fingerprint: portable_ref().params_fingerprint,
861        };
862        let handle = store.materialize(&request).unwrap();
863        assert_eq!(handle, 1);
864
865        let duplicate = store.register(portable_ref()).unwrap_err();
866        assert!(format!("{duplicate}").contains("already has handle for"));
867
868        let mismatch_request = FittedAdapterMaterializationRequest {
869            adapter_id: "snv".to_string(),
870            params_fingerprint: fingerprint(0x55),
871        };
872        let error = store.materialize(&mismatch_request).unwrap_err();
873        assert!(
874            format!("{error}").contains("params fingerprint mismatch"),
875            "unexpected: {error}"
876        );
877
878        let missing_request = FittedAdapterMaterializationRequest {
879            adapter_id: "msc".to_string(),
880            params_fingerprint: portable_ref().params_fingerprint,
881        };
882        let error = store.materialize(&missing_request).unwrap_err();
883        assert!(format!("{error}").contains("missing adapter"));
884
885        assert!(store.release("snv"));
886        assert!(store.is_empty());
887        assert!(store.materialize(&request).is_err());
888    }
889
890    #[test]
891    fn store_register_manifest_returns_handles_for_each_entry() {
892        let store = InMemoryFittedAdapterStore::new();
893        let mut second = portable_ref();
894        second.adapter_id = "msc".to_string();
895        let manifest = FittedAdapterManifest::new(vec![
896            FittedAdapterManifestEntry {
897                adapter_id: "snv".to_string(),
898                fitted_adapter: portable_ref(),
899            },
900            FittedAdapterManifestEntry {
901                adapter_id: "msc".to_string(),
902                fitted_adapter: second,
903            },
904        ]);
905        let records = store.register_manifest(&manifest).unwrap();
906        assert_eq!(records.len(), 2);
907        assert_eq!(records[0].handle, 1);
908        assert_eq!(records[1].handle, 2);
909        assert_eq!(store.len(), 2);
910    }
911
912    #[test]
913    fn store_materialize_rejects_invalid_request_fingerprint() {
914        let store = InMemoryFittedAdapterStore::new();
915        let request = FittedAdapterMaterializationRequest {
916            adapter_id: "snv".to_string(),
917            params_fingerprint: "not-a-fingerprint".to_string(),
918        };
919        let error = store.materialize(&request).unwrap_err();
920        assert!(format!("{error}").contains("params fingerprint"));
921    }
922
923    #[test]
924    fn manifest_portable_check_rejects_non_portable_entries() {
925        let entry = FittedAdapterManifestEntry {
926            adapter_id: "snv".to_string(),
927            fitted_adapter: FittedAdapterRef {
928                uri: None,
929                backend: None,
930                content_fingerprint: None,
931                size_bytes: None,
932                plugin: None,
933                plugin_version: None,
934                ..portable_ref()
935            },
936        };
937        let manifest = FittedAdapterManifest::new(vec![entry]);
938        manifest.validate().unwrap();
939        let error = manifest.validate_portable().unwrap_err();
940        assert!(format!("{error}").contains("is not portable"));
941    }
942}
dag_ml_data_core/fitted_adapter.rs

dag_ml_data_core/
fitted_adapter.rs
