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)]
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 Default for InMemoryFittedAdapterStore {
300    fn default() -> Self {
301        Self::new()
302    }
303}
304
305impl InMemoryFittedAdapterStore {
306    pub fn new() -> Self {
307        Self {
308            state: Mutex::new(FittedAdapterStoreState {
309                next_handle: 1,
310                records: BTreeMap::new(),
311            }),
312        }
313    }
314
315    /// Locks the store, propagating poison errors as a validation failure.
316    /// Poisoning happens when another thread panicked while holding the lock;
317    /// for an in-memory contract store the right policy is to surface the
318    /// error rather than silently `recover` from it.
319    fn lock(&self) -> Result<std::sync::MutexGuard<'_, FittedAdapterStoreState>> {
320        self.state.lock().map_err(|error| {
321            DataError::Validation(format!("fitted adapter store mutex poisoned: {error}"))
322        })
323    }
324
325    pub fn register(&self, fitted_adapter: FittedAdapterRef) -> Result<FittedAdapterHandleRecord> {
326        fitted_adapter.validate()?;
327        let mut state = self.lock()?;
328        if state.records.contains_key(&fitted_adapter.adapter_id) {
329            return Err(DataError::Validation(format!(
330                "fitted adapter store already has handle for `{}`",
331                fitted_adapter.adapter_id
332            )));
333        }
334        let handle = state.next_handle;
335        state.next_handle += 1;
336        let record = FittedAdapterHandleRecord {
337            handle,
338            fitted_adapter,
339        };
340        state
341            .records
342            .insert(record.fitted_adapter.adapter_id.clone(), record.clone());
343        Ok(record)
344    }
345
346    pub fn register_manifest(
347        &self,
348        manifest: &FittedAdapterManifest,
349    ) -> Result<Vec<FittedAdapterHandleRecord>> {
350        manifest.validate()?;
351        manifest
352            .entries
353            .iter()
354            .map(|entry| self.register(entry.fitted_adapter.clone()))
355            .collect()
356    }
357
358    pub fn get(&self, adapter_id: &str) -> Option<FittedAdapterHandleRecord> {
359        let state = self.lock().ok()?;
360        state.records.get(adapter_id).cloned()
361    }
362
363    pub fn release(&self, adapter_id: &str) -> bool {
364        match self.lock() {
365            Ok(mut state) => state.records.remove(adapter_id).is_some(),
366            Err(_) => false,
367        }
368    }
369
370    pub fn len(&self) -> usize {
371        self.lock().map(|state| state.records.len()).unwrap_or(0)
372    }
373
374    pub fn is_empty(&self) -> bool {
375        self.lock()
376            .map(|state| state.records.is_empty())
377            .unwrap_or(true)
378    }
379}
380
381impl RuntimeFittedAdapterStore for InMemoryFittedAdapterStore {
382    fn materialize(&self, request: &FittedAdapterMaterializationRequest) -> Result<u64> {
383        request.validate()?;
384        let state = self.lock()?;
385        let record = state.records.get(&request.adapter_id).ok_or_else(|| {
386            DataError::Validation(format!(
387                "fitted adapter store is missing adapter `{}`",
388                request.adapter_id
389            ))
390        })?;
391        if record.fitted_adapter.params_fingerprint != request.params_fingerprint {
392            return Err(DataError::FingerprintMismatch {
393                kind: "params",
394                expected: record.fitted_adapter.params_fingerprint.clone(),
395                actual: request.params_fingerprint.clone(),
396            });
397        }
398        Ok(record.handle)
399    }
400}
401
402fn validate_optional_text(label: &str, value: &Option<String>, adapter_id: &str) -> Result<()> {
403    if let Some(text) = value {
404        if text.trim().is_empty() {
405            return Err(DataError::Validation(format!(
406                "fitted adapter `{adapter_id}` has empty {label}"
407            )));
408        }
409        if text.chars().any(char::is_control) {
410            return Err(DataError::Validation(format!(
411                "fitted adapter `{adapter_id}` has control characters in {label}"
412            )));
413        }
414    }
415    Ok(())
416}
417
418fn validate_fingerprint(label: &str, value: &str, adapter_id: &str) -> Result<()> {
419    if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
420        return Err(DataError::Validation(format!(
421            "fitted adapter `{adapter_id}` {label} fingerprint must be a 64-character hex digest"
422        )));
423    }
424    Ok(())
425}
426
427/// Deterministic path safety for relative fitted-adapter URIs. Rejects empty
428/// values, control characters, absolute paths (POSIX root, Windows root or
429/// drive prefix), URI schemes such as `http://`, `s3://` or `file://` (any
430/// colon in the leading path segment) and any `..` traversal component.
431/// Parsing is platform-independent so portable manifests validate identically
432/// everywhere; it adds no dependency. Kept byte-for-byte equivalent with
433/// `dag-ml`'s `validate_relative_artifact_uri` in
434/// `crates/dag-ml-core/src/runtime/prediction_store.rs` so portable refs accepted by one repo
435/// are accepted by the other.
436fn validate_relative_uri(adapter_id: &str, uri: &str) -> Result<()> {
437    if uri.is_empty() {
438        return Err(DataError::Validation(format!(
439            "fitted adapter `{adapter_id}` has empty uri"
440        )));
441    }
442    if uri.chars().any(char::is_control) {
443        return Err(DataError::Validation(format!(
444            "fitted adapter `{adapter_id}` uri has control characters"
445        )));
446    }
447    if uri.starts_with('/') || uri.starts_with('\\') {
448        return Err(DataError::Validation(format!(
449            "fitted adapter `{adapter_id}` uri `{uri}` must be a relative path"
450        )));
451    }
452    let mut prefix = uri.chars();
453    if let (Some(drive), Some(':')) = (prefix.next(), prefix.next()) {
454        if drive.is_ascii_alphabetic() {
455            return Err(DataError::Validation(format!(
456                "fitted adapter `{adapter_id}` uri `{uri}` must be a relative path"
457            )));
458        }
459    }
460    let first_segment = uri.split(['/', '\\']).next().unwrap_or(uri);
461    if first_segment.contains(':') {
462        return Err(DataError::Validation(format!(
463            "fitted adapter `{adapter_id}` uri `{uri}` must not include a scheme or colon in its first path segment"
464        )));
465    }
466    for segment in uri.split(['/', '\\']) {
467        if segment == ".." {
468            return Err(DataError::Validation(format!(
469                "fitted adapter `{adapter_id}` uri `{uri}` must not contain `..` components"
470            )));
471        }
472    }
473    Ok(())
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479
480    fn fingerprint(byte: u8) -> String {
481        format!("{byte:02x}").repeat(32)
482    }
483
484    fn portable_ref() -> FittedAdapterRef {
485        FittedAdapterRef {
486            schema_version: FITTED_ADAPTER_REF_SCHEMA_VERSION,
487            adapter_id: "snv".to_string(),
488            adapter_version: "1.0.0".to_string(),
489            params_fingerprint: fingerprint(0x12),
490            backend: Some(FittedAdapterBackend::Joblib),
491            uri: Some("fitted/snv.joblib".to_string()),
492            content_fingerprint: Some(fingerprint(0x34)),
493            size_bytes: Some(2048),
494            plugin: Some("sklearn".to_string()),
495            plugin_version: Some("1.4.0".to_string()),
496            metadata: BTreeMap::new(),
497        }
498    }
499
500    #[test]
501    fn portable_ref_validates() {
502        let value = portable_ref();
503        value.validate().unwrap();
504        value.validate_portable().unwrap();
505    }
506
507    #[test]
508    fn inline_ref_validates_but_is_not_portable() {
509        let value = FittedAdapterRef {
510            uri: None,
511            backend: None,
512            content_fingerprint: None,
513            size_bytes: None,
514            plugin: None,
515            plugin_version: None,
516            ..portable_ref()
517        };
518        value.validate().unwrap();
519        let error = value.validate_portable().unwrap_err();
520        assert!(format!("{error}").contains("is not portable"));
521    }
522
523    #[test]
524    fn rejects_unsupported_schema_version() {
525        let mut value = portable_ref();
526        value.schema_version = FITTED_ADAPTER_REF_SCHEMA_VERSION + 1;
527        let error = value.validate().unwrap_err();
528        assert!(format!("{error}").contains("unsupported schema_version"));
529    }
530
531    #[test]
532    fn rejects_uri_without_backend_or_fingerprint() {
533        let mut value = portable_ref();
534        value.backend = None;
535        let error = value.validate().unwrap_err();
536        assert!(format!("{error}").contains("has uri without backend"));
537
538        let mut value = portable_ref();
539        value.content_fingerprint = None;
540        let error = value.validate().unwrap_err();
541        assert!(format!("{error}").contains("has uri without content_fingerprint"));
542    }
543
544    #[test]
545    fn rejects_plugin_version_without_plugin() {
546        let mut value = portable_ref();
547        value.plugin = None;
548        let error = value.validate().unwrap_err();
549        assert!(format!("{error}").contains("has plugin_version without plugin"));
550    }
551
552    /// Locks down the exact accept/reject behaviour of the portable URI
553    /// rules so the byte-for-byte copy in dag-ml's
554    /// `validate_relative_artifact_uri` can be cross-checked. Any change to
555    /// either function must also update this fixture; the script-level check
556    /// in `scripts/validate_contracts.py` reads both function bodies and
557    /// asserts textual equivalence as a second line of defense.
558    #[test]
559    fn portable_uri_rules_are_locked_to_a_shared_fixture() {
560        let accept = [
561            "snv.joblib",
562            "fitted/snv.joblib",
563            "nested/dir.name/file-v1.bin",
564            "alpha_beta/gamma.json",
565            "a.b/c.d",
566        ];
567        let reject_absolute = ["/abs/path.joblib", "\\abs\\path.joblib"];
568        let reject_drive = ["C:\\fitted\\snv.joblib", "Z:/path"];
569        let reject_schemes = [
570            "http://example.com/x.bin",
571            "s3://bucket/x.bin",
572            "file:///x.bin",
573            "scheme:relative",
574        ];
575        let reject_traversal = [
576            "../escape.bin",
577            "fitted/../escape.bin",
578            "ok/../../escape.bin",
579        ];
580        let reject_empty_or_control = ["", "ctrl/\u{0001}path", "\n/path"];
581
582        for uri in accept {
583            let value = FittedAdapterRef {
584                uri: Some(uri.to_string()),
585                ..portable_ref()
586            };
587            if let Err(error) = value.validate_portable() {
588                panic!("expected `{uri}` to validate portable: {error}");
589            }
590        }
591        for group in [
592            &reject_absolute[..],
593            &reject_drive[..],
594            &reject_schemes[..],
595            &reject_traversal[..],
596            &reject_empty_or_control[..],
597        ] {
598            for uri in group {
599                let value = FittedAdapterRef {
600                    uri: Some(uri.to_string()),
601                    ..portable_ref()
602                };
603                let error = value
604                    .validate_portable()
605                    .expect_err(&format!("expected `{uri}` to be rejected as non-portable"));
606                let message = format!("{error}");
607                assert!(
608                    !message.is_empty(),
609                    "rejection of `{uri}` must carry a non-empty error message"
610                );
611            }
612        }
613    }
614
615    #[test]
616    fn rejects_non_portable_uris() {
617        for uri in [
618            "/abs/path.joblib",
619            "C:\\fitted\\snv.joblib",
620            "http://example.com/fitted/snv.joblib",
621            "s3://bucket/fitted/snv.joblib",
622            "file:///fitted/snv.joblib",
623            "../escape/snv.joblib",
624            "fitted/../escape/snv.joblib",
625        ] {
626            let value = FittedAdapterRef {
627                uri: Some(uri.to_string()),
628                ..portable_ref()
629            };
630            let error = value.validate_portable().unwrap_err();
631            assert!(
632                format!("{error}").contains("must") || format!("{error}").contains("relative path"),
633                "URI `{uri}` should be rejected, got: {error}"
634            );
635        }
636    }
637
638    #[test]
639    fn rejects_uri_with_control_chars_or_empty() {
640        let mut value = portable_ref();
641        value.uri = Some("".to_string());
642        let error = value.validate_portable().unwrap_err();
643        assert!(format!("{error}").contains("empty uri"));
644
645        let mut value = portable_ref();
646        value.uri = Some("fitted/\u{0001}".to_string());
647        let error = value.validate_portable().unwrap_err();
648        assert!(format!("{error}").contains("control characters"));
649    }
650
651    #[test]
652    fn rejects_whitespace_only_optional_text_fields() {
653        for (label, value) in [
654            (
655                "uri",
656                FittedAdapterRef {
657                    uri: Some("   ".to_string()),
658                    ..portable_ref()
659                },
660            ),
661            (
662                "plugin",
663                FittedAdapterRef {
664                    plugin: Some("\t".to_string()),
665                    ..portable_ref()
666                },
667            ),
668            (
669                "plugin_version",
670                FittedAdapterRef {
671                    plugin_version: Some("   ".to_string()),
672                    ..portable_ref()
673                },
674            ),
675        ] {
676            let error = value.validate().unwrap_err();
677            let message = format!("{error}");
678            assert!(
679                message.contains(&format!("has empty {label}")),
680                "expected empty {label} error, got: {message}"
681            );
682        }
683    }
684
685    #[test]
686    fn published_fitted_adapter_ref_schema_pins_current_id_and_backends() {
687        let schema: serde_json::Value = serde_json::from_str(include_str!(
688            "../../../docs/contracts/fitted_adapter_ref.schema.json"
689        ))
690        .unwrap();
691        assert_eq!(schema["$id"].as_str(), Some(FITTED_ADAPTER_REF_SCHEMA_ID));
692        assert!(
693            schema["$id"]
694                .as_str()
695                .unwrap()
696                .ends_with(&format!("v{FITTED_ADAPTER_REF_SCHEMA_VERSION}.schema.json")),
697            "schema $id `{}` must encode version v{FITTED_ADAPTER_REF_SCHEMA_VERSION}",
698            schema["$id"]
699        );
700        let backends = schema["$defs"]["backend"]["enum"].as_array().unwrap();
701        for expected in ["joblib", "pickle", "json", "numpy", "onnx", "raw"] {
702            assert!(
703                backends
704                    .iter()
705                    .any(|value| value.as_str() == Some(expected)),
706                "fitted_adapter_ref backend enum is missing `{expected}`"
707            );
708        }
709    }
710
711    #[test]
712    fn rejects_metadata_with_empty_keys() {
713        let mut value = portable_ref();
714        value
715            .metadata
716            .insert("".to_string(), serde_json::Value::Null);
717        let error = value.validate().unwrap_err();
718        assert!(format!("{error}").contains("metadata contains an empty key"));
719    }
720
721    #[test]
722    fn rejects_invalid_fingerprints() {
723        let mut value = portable_ref();
724        value.params_fingerprint = "not-a-fingerprint".to_string();
725        let error = value.validate().unwrap_err();
726        assert!(format!("{error}").contains("params fingerprint"));
727
728        let mut value = portable_ref();
729        value.content_fingerprint = Some("not-a-fingerprint".to_string());
730        let error = value.validate().unwrap_err();
731        assert!(format!("{error}").contains("content fingerprint"));
732    }
733
734    #[test]
735    fn manifest_validates_and_requires_unique_adapter_ids() {
736        let entry = FittedAdapterManifestEntry {
737            adapter_id: "snv".to_string(),
738            fitted_adapter: portable_ref(),
739        };
740        let manifest = FittedAdapterManifest::new(vec![entry.clone()]);
741        manifest.validate().unwrap();
742        manifest.validate_portable().unwrap();
743
744        let manifest = FittedAdapterManifest::new(vec![entry.clone(), entry]);
745        let error = manifest.validate().unwrap_err();
746        assert!(format!("{error}").contains("duplicate adapter_id"));
747    }
748
749    #[test]
750    fn manifest_refuses_key_mismatch_and_bad_schema_version() {
751        let entry = FittedAdapterManifestEntry {
752            adapter_id: "snv".to_string(),
753            fitted_adapter: FittedAdapterRef {
754                adapter_id: "msc".to_string(),
755                ..portable_ref()
756            },
757        };
758        let manifest = FittedAdapterManifest::new(vec![entry]);
759        let error = manifest.validate().unwrap_err();
760        assert!(format!("{error}").contains("does not match ref adapter_id"));
761
762        let mut manifest = FittedAdapterManifest::new(vec![]);
763        manifest.schema_version = FITTED_ADAPTER_MANIFEST_SCHEMA_VERSION + 1;
764        let error = manifest.validate().unwrap_err();
765        assert!(format!("{error}").contains("unsupported schema_version"));
766    }
767
768    #[test]
769    fn manifest_round_trips_through_file_persistence() {
770        let manifest = FittedAdapterManifest::new(vec![FittedAdapterManifestEntry {
771            adapter_id: "snv".to_string(),
772            fitted_adapter: portable_ref(),
773        }]);
774        let path = std::env::temp_dir().join(format!(
775            "dag_ml_data_fitted_adapter_manifest_roundtrip_{}.json",
776            std::process::id()
777        ));
778        manifest.write_to_path(&path, true).unwrap();
779        let loaded = FittedAdapterManifest::read_from_path(&path, true).unwrap();
780        assert_eq!(loaded, manifest);
781        let _ = std::fs::remove_file(&path);
782    }
783
784    #[test]
785    fn manifest_file_persistence_rejects_inline_when_portable_required() {
786        let manifest = FittedAdapterManifest::new(vec![FittedAdapterManifestEntry {
787            adapter_id: "snv".to_string(),
788            fitted_adapter: FittedAdapterRef {
789                uri: None,
790                backend: None,
791                content_fingerprint: None,
792                plugin: None,
793                plugin_version: None,
794                ..portable_ref()
795            },
796        }]);
797        let path = std::env::temp_dir().join(format!(
798            "dag_ml_data_fitted_adapter_manifest_inline_{}.json",
799            std::process::id()
800        ));
801        let error = manifest.write_to_path(&path, true).unwrap_err();
802        assert!(format!("{error}").contains("is not portable"));
803        manifest.write_to_path(&path, false).unwrap();
804        let loaded_inline = FittedAdapterManifest::read_from_path(&path, false).unwrap();
805        assert_eq!(loaded_inline, manifest);
806        let portable_load_error = FittedAdapterManifest::read_from_path(&path, true).unwrap_err();
807        assert!(format!("{portable_load_error}").contains("is not portable"));
808        let _ = std::fs::remove_file(&path);
809    }
810
811    #[test]
812    fn manifest_file_persistence_reports_missing_paths() {
813        let path = std::env::temp_dir().join(format!(
814            "dag_ml_data_fitted_adapter_manifest_missing_{}.json",
815            std::process::id()
816        ));
817        // Ensure path does not exist.
818        let _ = std::fs::remove_file(&path);
819        let error = FittedAdapterManifest::read_from_path(&path, false).unwrap_err();
820        assert!(format!("{error}").contains("failed to read fitted adapter manifest"));
821    }
822
823    #[test]
824    fn in_memory_store_is_sync_and_handles_concurrent_register_calls() {
825        use std::sync::Arc;
826
827        // Compile-time check: store must be Send + Sync so it can be shared
828        // across host threads through the C ABI handle.
829        fn assert_send_sync<T: Send + Sync>() {}
830        assert_send_sync::<InMemoryFittedAdapterStore>();
831
832        let store = Arc::new(InMemoryFittedAdapterStore::new());
833        let mut handles = Vec::new();
834        for thread_idx in 0..8 {
835            let store = store.clone();
836            handles.push(std::thread::spawn(move || {
837                let mut record = portable_ref();
838                record.adapter_id = format!("snv:{thread_idx}");
839                store.register(record).unwrap()
840            }));
841        }
842        let records: Vec<_> = handles
843            .into_iter()
844            .map(|handle| handle.join().unwrap())
845            .collect();
846        assert_eq!(records.len(), 8);
847        assert_eq!(store.len(), 8);
848        let mut handle_ids: Vec<u64> = records.iter().map(|record| record.handle).collect();
849        handle_ids.sort();
850        // Handles are assigned 1..=8 across the threads, no duplicates.
851        assert_eq!(handle_ids, (1..=8).collect::<Vec<u64>>());
852    }
853
854    #[test]
855    fn in_memory_store_registers_and_materializes_adapter_handles() {
856        let store = InMemoryFittedAdapterStore::new();
857        assert!(store.is_empty());
858
859        let record = store.register(portable_ref()).unwrap();
860        assert_eq!(record.handle, 1);
861        assert_eq!(record.fitted_adapter.adapter_id, "snv");
862        assert_eq!(store.len(), 1);
863
864        let request = FittedAdapterMaterializationRequest {
865            adapter_id: "snv".to_string(),
866            params_fingerprint: portable_ref().params_fingerprint,
867        };
868        let handle = store.materialize(&request).unwrap();
869        assert_eq!(handle, 1);
870
871        let duplicate = store.register(portable_ref()).unwrap_err();
872        assert!(format!("{duplicate}").contains("already has handle for"));
873
874        let mismatch_request = FittedAdapterMaterializationRequest {
875            adapter_id: "snv".to_string(),
876            params_fingerprint: fingerprint(0x55),
877        };
878        let error = store.materialize(&mismatch_request).unwrap_err();
879        assert!(
880            format!("{error}").contains("params fingerprint mismatch"),
881            "unexpected: {error}"
882        );
883
884        let missing_request = FittedAdapterMaterializationRequest {
885            adapter_id: "msc".to_string(),
886            params_fingerprint: portable_ref().params_fingerprint,
887        };
888        let error = store.materialize(&missing_request).unwrap_err();
889        assert!(format!("{error}").contains("missing adapter"));
890
891        assert!(store.release("snv"));
892        assert!(store.is_empty());
893        assert!(store.materialize(&request).is_err());
894    }
895
896    #[test]
897    fn store_register_manifest_returns_handles_for_each_entry() {
898        let store = InMemoryFittedAdapterStore::new();
899        let mut second = portable_ref();
900        second.adapter_id = "msc".to_string();
901        let manifest = FittedAdapterManifest::new(vec![
902            FittedAdapterManifestEntry {
903                adapter_id: "snv".to_string(),
904                fitted_adapter: portable_ref(),
905            },
906            FittedAdapterManifestEntry {
907                adapter_id: "msc".to_string(),
908                fitted_adapter: second,
909            },
910        ]);
911        let records = store.register_manifest(&manifest).unwrap();
912        assert_eq!(records.len(), 2);
913        assert_eq!(records[0].handle, 1);
914        assert_eq!(records[1].handle, 2);
915        assert_eq!(store.len(), 2);
916    }
917
918    #[test]
919    fn store_materialize_rejects_invalid_request_fingerprint() {
920        let store = InMemoryFittedAdapterStore::new();
921        let request = FittedAdapterMaterializationRequest {
922            adapter_id: "snv".to_string(),
923            params_fingerprint: "not-a-fingerprint".to_string(),
924        };
925        let error = store.materialize(&request).unwrap_err();
926        assert!(format!("{error}").contains("params fingerprint"));
927    }
928
929    #[test]
930    fn manifest_portable_check_rejects_non_portable_entries() {
931        let entry = FittedAdapterManifestEntry {
932            adapter_id: "snv".to_string(),
933            fitted_adapter: FittedAdapterRef {
934                uri: None,
935                backend: None,
936                content_fingerprint: None,
937                size_bytes: None,
938                plugin: None,
939                plugin_version: None,
940                ..portable_ref()
941            },
942        };
943        let manifest = FittedAdapterManifest::new(vec![entry]);
944        manifest.validate().unwrap();
945        let error = manifest.validate_portable().unwrap_err();
946        assert!(format!("{error}").contains("is not portable"));
947    }
948}
dag_ml_data_core/fitted_adapter.rs

dag_ml_data_core/
fitted_adapter.rs
