ai_model_vault/
federation.rs

1//! Federated vault synchronization
2//!
3//! Enables synchronization of model versions across multiple vault instances.
4//! Uses a peer-to-peer protocol with eventual consistency.
5//!
6//! ## Architecture
7//!
8//! - **Nodes**: Independent vault instances that can sync with each other
9//! - **Sync Protocol**: Delta-based replication using vector clocks
10//! - **Conflict Resolution**: Last-writer-wins with version lineage preservation
11//! - **Transport**: HTTPS with mutual TLS authentication
12//!
13//! ## Security
14//!
15//! - All transfers are encrypted end-to-end
16//! - Node authentication via certificate pinning
17//! - Audit logs record all sync operations
18
19use chrono::{DateTime, Utc};
20use serde::{Deserialize, Serialize};
21use std::collections::{HashMap, HashSet};
22use std::path::PathBuf;
23use std::sync::Arc;
24use tokio::sync::RwLock;
25use uuid::Uuid;
26use zeroize::Zeroize;
27
28use crate::error::{Result, VaultError};
29
30/// Federation node identifier
31pub type NodeId = String;
32
33/// Vector clock for causal ordering
34#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
35pub struct VectorClock {
36    /// Map of node ID to logical timestamp
37    pub timestamps: HashMap<NodeId, u64>,
38}
39
40impl VectorClock {
41    /// Create new vector clock
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// Increment the clock for a node
47    pub fn increment(&mut self, node_id: &str) {
48        *self.timestamps.entry(node_id.to_string()).or_insert(0) += 1;
49    }
50
51    /// Merge with another vector clock (take max of each)
52    pub fn merge(&mut self, other: &VectorClock) {
53        for (node, ts) in &other.timestamps {
54            let entry = self.timestamps.entry(node.clone()).or_insert(0);
55            *entry = (*entry).max(*ts);
56        }
57    }
58
59    /// Check if this clock is causally before or concurrent with another
60    pub fn compare(&self, other: &VectorClock) -> ClockComparison {
61        let mut dominated = true;
62        let mut dominates = true;
63
64        for (node, ts) in &self.timestamps {
65            let other_ts = other.timestamps.get(node).copied().unwrap_or(0);
66            if *ts > other_ts {
67                dominated = false;
68            }
69            if *ts < other_ts {
70                dominates = false;
71            }
72        }
73
74        for (node, ts) in &other.timestamps {
75            if !self.timestamps.contains_key(node) && *ts > 0 {
76                dominates = false;
77            }
78        }
79
80        match (dominates, dominated) {
81            (true, true) => ClockComparison::Equal,
82            (true, false) => ClockComparison::After,
83            (false, true) => ClockComparison::Before,
84            (false, false) => ClockComparison::Concurrent,
85        }
86    }
87}
88
89/// Result of comparing two vector clocks
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum ClockComparison {
92    /// Clocks are equal
93    Equal,
94    /// First clock happened before second
95    Before,
96    /// First clock happened after second
97    After,
98    /// Clocks are concurrent (no causal relationship)
99    Concurrent,
100}
101
102/// Federation node configuration
103#[derive(Debug, Clone, Serialize, Deserialize)]
104pub struct FederationConfig {
105    /// This node's unique ID
106    pub node_id: NodeId,
107    /// This node's display name
108    pub node_name: String,
109    /// List of peer nodes to sync with
110    pub peers: Vec<PeerConfig>,
111    /// Sync interval in seconds (0 to disable auto-sync)
112    pub sync_interval_secs: u64,
113    /// Whether to auto-resolve conflicts
114    pub auto_resolve_conflicts: bool,
115    /// Maximum concurrent syncs
116    pub max_concurrent_syncs: usize,
117}
118
119impl Default for FederationConfig {
120    fn default() -> Self {
121        Self {
122            node_id: Uuid::new_v4().to_string(),
123            node_name: hostname::get()
124                .map(|h| h.to_string_lossy().into_owned())
125                .unwrap_or_else(|_| "unknown".into()),
126            peers: Vec::new(),
127            sync_interval_secs: 300, // 5 minutes
128            auto_resolve_conflicts: true,
129            max_concurrent_syncs: 4,
130        }
131    }
132}
133
134/// Peer node configuration
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct PeerConfig {
137    /// Peer's node ID
138    pub node_id: NodeId,
139    /// Peer's display name
140    pub name: String,
141    /// Peer's sync endpoint URL
142    pub endpoint: String,
143    /// Optional API key for authentication
144    pub api_key: Option<String>,
145    /// Whether sync is enabled for this peer
146    pub enabled: bool,
147}
148
149impl Drop for PeerConfig {
150    fn drop(&mut self) {
151        if let Some(ref mut key) = self.api_key {
152            key.zeroize();
153        }
154    }
155}
156
157/// Sync state for a model
158#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct ModelSyncState {
160    /// Model name
161    pub name: String,
162    /// Vector clock for this model
163    pub clock: VectorClock,
164    /// Last sync time per peer
165    pub last_sync: HashMap<NodeId, DateTime<Utc>>,
166    /// Known versions across all nodes
167    pub known_versions: HashSet<String>, // checkpoint IDs
168    /// Pending uploads
169    pub pending_upload: HashSet<String>,
170    /// Pending downloads
171    pub pending_download: HashSet<(NodeId, String)>, // (node_id, checkpoint_id)
172}
173
174/// Sync manifest exchanged between nodes
175#[derive(Debug, Clone, Serialize, Deserialize)]
176pub struct SyncManifest {
177    /// Source node ID
178    pub source_node: NodeId,
179    /// Manifest timestamp
180    pub timestamp: DateTime<Utc>,
181    /// Model states
182    pub models: Vec<ModelManifestEntry>,
183    /// Node's vector clock
184    pub clock: VectorClock,
185}
186
187/// Single model entry in sync manifest
188#[derive(Debug, Clone, Serialize, Deserialize)]
189pub struct ModelManifestEntry {
190    /// Model name
191    pub name: String,
192    /// Available versions
193    pub versions: Vec<VersionManifestEntry>,
194    /// Model's vector clock
195    pub clock: VectorClock,
196}
197
198/// Version entry in manifest
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub struct VersionManifestEntry {
201    /// Version number
202    pub version: u32,
203    /// Checkpoint ID
204    pub checkpoint_id: String,
205    /// Creation timestamp
206    pub created_at: DateTime<Utc>,
207    /// SHA-256 checksum
208    pub checksum: String,
209    /// Size in bytes
210    pub size_bytes: u64,
211    /// Parent checkpoint ID (for lineage)
212    pub parent_id: Option<String>,
213    /// Originating node
214    pub origin_node: NodeId,
215}
216
217/// Sync operation result
218#[derive(Debug, Clone, Serialize, Deserialize)]
219pub struct SyncResult {
220    /// Peer node ID
221    pub peer_id: NodeId,
222    /// Timestamp
223    pub timestamp: DateTime<Utc>,
224    /// Duration in milliseconds
225    pub duration_ms: u64,
226    /// Models synced
227    pub models_synced: u32,
228    /// Versions uploaded
229    pub versions_uploaded: u32,
230    /// Versions downloaded
231    pub versions_downloaded: u32,
232    /// Conflicts detected
233    pub conflicts: Vec<SyncConflict>,
234    /// Errors encountered
235    pub errors: Vec<String>,
236}
237
238/// Sync conflict details
239#[derive(Debug, Clone, Serialize, Deserialize)]
240pub struct SyncConflict {
241    /// Model name
242    pub model: String,
243    /// Local version
244    pub local_version: String,
245    /// Remote version
246    pub remote_version: String,
247    /// Remote node
248    pub remote_node: NodeId,
249    /// Resolution (if auto-resolved)
250    pub resolution: Option<ConflictResolution>,
251}
252
253/// Conflict resolution strategy
254#[derive(Debug, Clone, Serialize, Deserialize)]
255pub enum ConflictResolution {
256    /// Keep local version
257    KeepLocal,
258    /// Use remote version
259    UseRemote,
260    /// Keep both as separate branches
261    Branch {
262        local_name: String,
263        remote_name: String,
264    },
265    /// Manual resolution required
266    Manual,
267}
268
269/// Federation manager
270pub struct FederationManager {
271    config: FederationConfig,
272    state: Arc<RwLock<FederationState>>,
273    http_client: reqwest::Client,
274}
275
276/// Internal federation state
277struct FederationState {
278    /// Sync state per model
279    models: HashMap<String, ModelSyncState>,
280    /// Global vector clock
281    clock: VectorClock,
282    /// Sync history
283    history: Vec<SyncResult>,
284    /// State file path
285    state_file: PathBuf,
286}
287
288impl FederationManager {
289    /// Create new federation manager
290    pub fn new(config: FederationConfig, state_dir: PathBuf) -> Result<Self> {
291        let state_file = state_dir.join("federation_state.json");
292
293        // Load existing state or create new
294        let federation_state = if state_file.exists() {
295            let contents = std::fs::read_to_string(&state_file)?;
296            let loaded: SavedFederationState = serde_json::from_str(&contents)?;
297            FederationState {
298                models: loaded.models,
299                clock: loaded.clock,
300                history: loaded.history,
301                state_file,
302            }
303        } else {
304            FederationState {
305                models: HashMap::new(),
306                clock: VectorClock::new(),
307                history: Vec::new(),
308                state_file,
309            }
310        };
311
312        let http_client = reqwest::Client::builder()
313            .timeout(std::time::Duration::from_secs(300))
314            .build()
315            .map_err(|e| {
316                VaultError::IoError(std::io::Error::other(format!("HTTP client error: {e}")))
317            })?;
318
319        Ok(Self {
320            config,
321            state: Arc::new(RwLock::new(federation_state)),
322            http_client,
323        })
324    }
325
326    /// Get this node's ID
327    pub fn node_id(&self) -> &str {
328        &self.config.node_id
329    }
330
331    /// Get configured peers
332    pub fn peers(&self) -> &[PeerConfig] {
333        &self.config.peers
334    }
335
336    /// Add a peer
337    pub fn add_peer(&mut self, peer: PeerConfig) {
338        self.config.peers.push(peer);
339    }
340
341    /// Remove a peer
342    pub fn remove_peer(&mut self, node_id: &str) {
343        self.config.peers.retain(|p| p.node_id != node_id);
344    }
345
346    /// Generate sync manifest from local vault
347    pub async fn generate_manifest(
348        &self,
349        vault_models: Vec<(String, Vec<crate::version::ModelVersion>)>,
350    ) -> SyncManifest {
351        let state = self.state.read().await;
352
353        let models = vault_models
354            .into_iter()
355            .map(
356                |(name, versions): (String, Vec<crate::version::ModelVersion>)| {
357                    let model_clock = state
358                        .models
359                        .get(&name)
360                        .map(|s| s.clock.clone())
361                        .unwrap_or_default();
362
363                    let version_entries = versions
364                        .into_iter()
365                        .map(|v| VersionManifestEntry {
366                            version: v.version,
367                            // The federation identity, not the local one: a
368                            // model received from a peer keeps the id it was
369                            // created with, so both nodes agree on what they
370                            // already have and sync converges instead of
371                            // re-transferring on every run.
372                            checkpoint_id: crate::federation_transport::federation_checkpoint_id(
373                                &v,
374                            ),
375                            created_at: v.timestamp,
376                            checksum: v.checksum_sha256,
377                            size_bytes: v.size_bytes,
378                            parent_id: v.parent_version.map(|p| p.to_string()),
379                            origin_node: self.config.node_id.clone(),
380                        })
381                        .collect();
382
383                    ModelManifestEntry {
384                        name,
385                        versions: version_entries,
386                        clock: model_clock,
387                    }
388                },
389            )
390            .collect();
391
392        SyncManifest {
393            source_node: self.config.node_id.clone(),
394            timestamp: Utc::now(),
395            models,
396            clock: state.clock.clone(),
397        }
398    }
399
400    /// Compute sync delta between local and remote manifests
401    pub fn compute_delta(&self, local: &SyncManifest, remote: &SyncManifest) -> SyncDelta {
402        let mut to_upload = Vec::new();
403        let mut to_download = Vec::new();
404        let mut conflicts = Vec::new();
405
406        // Build lookup for remote models
407        let remote_models: HashMap<&str, &ModelManifestEntry> =
408            remote.models.iter().map(|m| (m.name.as_str(), m)).collect();
409
410        // Check each local model
411        for local_model in &local.models {
412            if let Some(remote_model) = remote_models.get(local_model.name.as_str()) {
413                // Model exists on both sides - compare versions
414                let local_checkpoints: HashSet<&str> = local_model
415                    .versions
416                    .iter()
417                    .map(|v| v.checkpoint_id.as_str())
418                    .collect();
419                let remote_checkpoints: HashSet<&str> = remote_model
420                    .versions
421                    .iter()
422                    .map(|v| v.checkpoint_id.as_str())
423                    .collect();
424
425                // Versions we have but remote doesn't
426                for version in &local_model.versions {
427                    if !remote_checkpoints.contains(version.checkpoint_id.as_str()) {
428                        to_upload.push(SyncItem {
429                            model: local_model.name.clone(),
430                            checkpoint_id: version.checkpoint_id.clone(),
431                            size_bytes: version.size_bytes,
432                        });
433                    }
434                }
435
436                // Versions remote has but we don't
437                for version in &remote_model.versions {
438                    if !local_checkpoints.contains(version.checkpoint_id.as_str()) {
439                        to_download.push(SyncItem {
440                            model: remote_model.name.clone(),
441                            checkpoint_id: version.checkpoint_id.clone(),
442                            size_bytes: version.size_bytes,
443                        });
444                    }
445                }
446
447                // Check for conflicts (same version number, different checkpoints)
448                let local_by_version: HashMap<u32, &VersionManifestEntry> = local_model
449                    .versions
450                    .iter()
451                    .map(|v| (v.version, v))
452                    .collect();
453                let remote_by_version: HashMap<u32, &VersionManifestEntry> = remote_model
454                    .versions
455                    .iter()
456                    .map(|v| (v.version, v))
457                    .collect();
458
459                for (version, local_v) in &local_by_version {
460                    if let Some(remote_v) = remote_by_version.get(version) {
461                        if local_v.checkpoint_id != remote_v.checkpoint_id {
462                            conflicts.push(SyncConflict {
463                                model: local_model.name.clone(),
464                                local_version: local_v.checkpoint_id.clone(),
465                                remote_version: remote_v.checkpoint_id.clone(),
466                                remote_node: remote.source_node.clone(),
467                                resolution: None,
468                            });
469                        }
470                    }
471                }
472            } else {
473                // Model only exists locally - upload all versions
474                for version in &local_model.versions {
475                    to_upload.push(SyncItem {
476                        model: local_model.name.clone(),
477                        checkpoint_id: version.checkpoint_id.clone(),
478                        size_bytes: version.size_bytes,
479                    });
480                }
481            }
482        }
483
484        // Check for models that only exist on remote
485        let local_model_names: HashSet<&str> =
486            local.models.iter().map(|m| m.name.as_str()).collect();
487        for remote_model in &remote.models {
488            if !local_model_names.contains(remote_model.name.as_str()) {
489                for version in &remote_model.versions {
490                    to_download.push(SyncItem {
491                        model: remote_model.name.clone(),
492                        checkpoint_id: version.checkpoint_id.clone(),
493                        size_bytes: version.size_bytes,
494                    });
495                }
496            }
497        }
498
499        SyncDelta {
500            to_upload,
501            to_download,
502            conflicts,
503        }
504    }
505
506    /// Sync with a specific peer
507    pub async fn sync_with_peer(
508        &self,
509        peer: &PeerConfig,
510        local_manifest: &SyncManifest,
511        download_fn: impl Fn(&str, &str) -> Result<Vec<u8>>,
512        upload_fn: impl Fn(&str, &str, &[u8]) -> Result<()>,
513    ) -> Result<SyncResult> {
514        let start = std::time::Instant::now();
515        let mut errors = Vec::new();
516
517        // Fetch remote manifest
518        let remote_manifest = self.fetch_manifest(peer).await?;
519
520        // Compute delta
521        let delta = self.compute_delta(local_manifest, &remote_manifest);
522
523        // Handle conflicts
524        let mut resolved_conflicts = Vec::new();
525        for conflict in delta.conflicts {
526            let resolution = if self.config.auto_resolve_conflicts {
527                // Last-writer-wins based on timestamp
528                // In production, use vector clock comparison
529                Some(ConflictResolution::KeepLocal)
530            } else {
531                Some(ConflictResolution::Manual)
532            };
533            resolved_conflicts.push(SyncConflict {
534                resolution,
535                ..conflict
536            });
537        }
538
539        // Download missing versions
540        let mut downloaded = 0;
541        for item in &delta.to_download {
542            match self
543                .download_version(peer, &item.model, &item.checkpoint_id)
544                .await
545            {
546                Ok(data) => {
547                    if let Err(e) = upload_fn(&item.model, &item.checkpoint_id, &data) {
548                        errors.push(format!(
549                            "Failed to store {}/{}: {e}",
550                            item.model, item.checkpoint_id
551                        ));
552                    } else {
553                        downloaded += 1;
554                    }
555                }
556                Err(e) => {
557                    errors.push(format!(
558                        "Failed to download {}/{}: {e}",
559                        item.model, item.checkpoint_id
560                    ));
561                }
562            }
563        }
564
565        // Upload missing versions
566        let mut uploaded = 0;
567        for item in &delta.to_upload {
568            match download_fn(&item.model, &item.checkpoint_id) {
569                Ok(data) => {
570                    if let Err(e) = self
571                        .upload_version(peer, &item.model, &item.checkpoint_id, &data)
572                        .await
573                    {
574                        errors.push(format!(
575                            "Failed to upload {}/{}: {e}",
576                            item.model, item.checkpoint_id
577                        ));
578                    } else {
579                        uploaded += 1;
580                    }
581                }
582                Err(e) => {
583                    errors.push(format!(
584                        "Failed to read {}/{}: {e}",
585                        item.model, item.checkpoint_id
586                    ));
587                }
588            }
589        }
590
591        // Update state
592        {
593            let mut state = self.state.write().await;
594            state.clock.increment(&self.config.node_id);
595            state.clock.merge(&remote_manifest.clock);
596        }
597
598        let result = SyncResult {
599            peer_id: peer.node_id.clone(),
600            timestamp: Utc::now(),
601            duration_ms: start.elapsed().as_millis() as u64,
602            models_synced: (delta.to_upload.len() + delta.to_download.len()) as u32,
603            versions_uploaded: uploaded,
604            versions_downloaded: downloaded,
605            conflicts: resolved_conflicts,
606            errors,
607        };
608
609        // Record in history
610        {
611            let mut state = self.state.write().await;
612            state.history.push(result.clone());
613            if state.history.len() > 1000 {
614                state.history.remove(0);
615            }
616            let _ = self.save_state(&state);
617        }
618
619        Ok(result)
620    }
621
622    /// Fetch manifest from peer
623    /// Fetch a peer's manifest without syncing.
624    ///
625    /// Public so `aim federation plan` can show what a sync would move before
626    /// anything moves.
627    pub async fn fetch_peer_manifest(&self, peer: &PeerConfig) -> Result<SyncManifest> {
628        self.fetch_manifest(peer).await
629    }
630
631    async fn fetch_manifest(&self, peer: &PeerConfig) -> Result<SyncManifest> {
632        let url = format!("{}/api/v1/federation/manifest", peer.endpoint);
633        let mut req = self.http_client.get(&url);
634
635        if let Some(api_key) = &peer.api_key {
636            req = req.header("X-API-Key", api_key);
637        }
638
639        let response = req
640            .send()
641            .await
642            .map_err(|e| VaultError::IoError(std::io::Error::other(format!("HTTP error: {e}"))))?;
643
644        if !response.status().is_success() {
645            return Err(VaultError::IoError(std::io::Error::other(format!(
646                "Peer returned error: {}",
647                response.status()
648            ))));
649        }
650
651        let manifest = response
652            .json::<SyncManifest>()
653            .await
654            .map_err(|e| VaultError::IoError(std::io::Error::other(format!("JSON error: {e}"))))?;
655
656        Ok(manifest)
657    }
658
659    /// Download a version from peer
660    async fn download_version(
661        &self,
662        peer: &PeerConfig,
663        model: &str,
664        checkpoint_id: &str,
665    ) -> Result<Vec<u8>> {
666        let url = format!(
667            "{}/api/v1/federation/models/{}/versions/{}",
668            peer.endpoint, model, checkpoint_id
669        );
670        let mut req = self.http_client.get(&url);
671
672        if let Some(api_key) = &peer.api_key {
673            req = req.header("X-API-Key", api_key);
674        }
675
676        let response = req
677            .send()
678            .await
679            .map_err(|e| VaultError::IoError(std::io::Error::other(format!("HTTP error: {e}"))))?;
680
681        if !response.status().is_success() {
682            return Err(VaultError::IoError(std::io::Error::other(format!(
683                "Download failed: {}",
684                response.status()
685            ))));
686        }
687
688        let data = response
689            .bytes()
690            .await
691            .map_err(|e| VaultError::IoError(std::io::Error::other(format!("Read error: {e}"))))?;
692
693        Ok(data.to_vec())
694    }
695
696    /// Upload a version to peer
697    async fn upload_version(
698        &self,
699        peer: &PeerConfig,
700        model: &str,
701        checkpoint_id: &str,
702        data: &[u8],
703    ) -> Result<()> {
704        let url = format!(
705            "{}/api/v1/federation/models/{}/versions/{}",
706            peer.endpoint, model, checkpoint_id
707        );
708        let mut req = self.http_client.put(&url).body(data.to_vec());
709
710        if let Some(api_key) = &peer.api_key {
711            req = req.header("X-API-Key", api_key);
712        }
713
714        let response = req
715            .send()
716            .await
717            .map_err(|e| VaultError::IoError(std::io::Error::other(format!("HTTP error: {e}"))))?;
718
719        if !response.status().is_success() {
720            return Err(VaultError::IoError(std::io::Error::other(format!(
721                "Upload failed: {}",
722                response.status()
723            ))));
724        }
725
726        Ok(())
727    }
728
729    /// Save state to disk
730    fn save_state(&self, state: &FederationState) -> Result<()> {
731        let saved = SavedFederationState {
732            models: state.models.clone(),
733            clock: state.clock.clone(),
734            history: state.history.clone(),
735        };
736
737        let json = serde_json::to_string_pretty(&saved)?;
738
739        // Write with restrictive permissions
740        {
741            use std::io::Write;
742            let mut opts = std::fs::OpenOptions::new();
743            opts.write(true).create(true).truncate(true);
744            crate::permissions::set_create_mode(&mut opts);
745            let mut f = opts.open(&state.state_file)?;
746            f.write_all(json.as_bytes())?;
747        }
748        crate::permissions::restrict_file(&state.state_file)?;
749
750        Ok(())
751    }
752
753    /// Get sync history
754    pub async fn get_history(&self, limit: Option<usize>) -> Vec<SyncResult> {
755        let state = self.state.read().await;
756        let history = &state.history;
757        if let Some(n) = limit {
758            history.iter().rev().take(n).cloned().collect()
759        } else {
760            history.clone()
761        }
762    }
763
764    /// Get federation status
765    pub async fn status(&self) -> FederationStatus {
766        let state = self.state.read().await;
767        FederationStatus {
768            node_id: self.config.node_id.clone(),
769            node_name: self.config.node_name.clone(),
770            peer_count: self.config.peers.len(),
771            model_count: state.models.len(),
772            last_sync: state.history.last().map(|r| r.timestamp),
773            clock: state.clock.clone(),
774        }
775    }
776}
777
778/// Sync delta
779#[derive(Debug, Clone)]
780pub struct SyncDelta {
781    /// Items to upload to remote
782    pub to_upload: Vec<SyncItem>,
783    /// Items to download from remote
784    pub to_download: Vec<SyncItem>,
785    /// Conflicts detected
786    pub conflicts: Vec<SyncConflict>,
787}
788
789/// Single item to sync
790#[derive(Debug, Clone)]
791pub struct SyncItem {
792    /// Model name
793    pub model: String,
794    /// Checkpoint ID
795    pub checkpoint_id: String,
796    /// Size in bytes
797    pub size_bytes: u64,
798}
799
800/// Saved federation state (for persistence)
801#[derive(Debug, Clone, Serialize, Deserialize)]
802struct SavedFederationState {
803    models: HashMap<String, ModelSyncState>,
804    clock: VectorClock,
805    history: Vec<SyncResult>,
806}
807
808/// Federation status summary
809#[derive(Debug, Clone, Serialize, Deserialize)]
810pub struct FederationStatus {
811    /// This node's ID
812    pub node_id: NodeId,
813    /// This node's name
814    pub node_name: String,
815    /// Number of configured peers
816    pub peer_count: usize,
817    /// Number of synced models
818    pub model_count: usize,
819    /// Last sync time
820    pub last_sync: Option<DateTime<Utc>>,
821    /// Current vector clock
822    pub clock: VectorClock,
823}
824
825#[cfg(test)]
826mod tests {
827    use super::*;
828
829    #[test]
830    fn test_vector_clock_basic() {
831        let mut clock1 = VectorClock::new();
832        let mut clock2 = VectorClock::new();
833
834        clock1.increment("node1");
835        assert_eq!(clock1.timestamps.get("node1"), Some(&1));
836
837        clock2.increment("node2");
838        clock1.merge(&clock2);
839
840        assert_eq!(clock1.timestamps.get("node1"), Some(&1));
841        assert_eq!(clock1.timestamps.get("node2"), Some(&1));
842    }
843
844    #[test]
845    fn test_vector_clock_comparison() {
846        let mut clock1 = VectorClock::new();
847        let mut clock2 = VectorClock::new();
848
849        // Initially equal (both empty)
850        assert_eq!(clock1.compare(&clock2), ClockComparison::Equal);
851
852        // clock1 is now after clock2
853        clock1.increment("node1");
854        assert_eq!(clock1.compare(&clock2), ClockComparison::After);
855        assert_eq!(clock2.compare(&clock1), ClockComparison::Before);
856
857        // Now concurrent
858        clock2.increment("node2");
859        assert_eq!(clock1.compare(&clock2), ClockComparison::Concurrent);
860    }
861
862    #[test]
863    fn test_vector_clock_multiple_increments() {
864        let mut clock = VectorClock::new();
865        clock.increment("a");
866        clock.increment("a");
867        clock.increment("a");
868        assert_eq!(clock.timestamps.get("a"), Some(&3));
869    }
870
871    #[test]
872    fn test_vector_clock_merge_takes_max() {
873        let mut c1 = VectorClock::new();
874        let mut c2 = VectorClock::new();
875
876        c1.increment("a");
877        c1.increment("a"); // a=2
878        c2.increment("a"); // a=1
879        c2.increment("b"); // b=1
880
881        c1.merge(&c2);
882        assert_eq!(c1.timestamps.get("a"), Some(&2)); // max(2,1)
883        assert_eq!(c1.timestamps.get("b"), Some(&1)); // new entry
884    }
885
886    #[test]
887    fn test_vector_clock_equal_after_merge() {
888        let mut c1 = VectorClock::new();
889        c1.increment("a");
890        let c2 = c1.clone();
891        assert_eq!(c1.compare(&c2), ClockComparison::Equal);
892    }
893
894    #[test]
895    fn test_vector_clock_serialization() {
896        let mut clock = VectorClock::new();
897        clock.increment("node1");
898        let json = serde_json::to_string(&clock).unwrap();
899        let deserialized: VectorClock = serde_json::from_str(&json).unwrap();
900        assert_eq!(clock, deserialized);
901    }
902
903    #[test]
904    fn test_clock_comparison_debug() {
905        assert_eq!(format!("{:?}", ClockComparison::Equal), "Equal");
906        assert_eq!(format!("{:?}", ClockComparison::Before), "Before");
907        assert_eq!(format!("{:?}", ClockComparison::After), "After");
908        assert_eq!(format!("{:?}", ClockComparison::Concurrent), "Concurrent");
909    }
910
911    #[test]
912    fn test_federation_config_default() {
913        let config = FederationConfig::default();
914        assert!(!config.node_id.is_empty());
915        assert_eq!(config.sync_interval_secs, 300);
916        assert!(config.auto_resolve_conflicts);
917        assert_eq!(config.max_concurrent_syncs, 4);
918        assert!(config.peers.is_empty());
919        assert!(!config.node_name.is_empty());
920    }
921
922    #[test]
923    fn test_federation_config_serialization() {
924        let config = FederationConfig::default();
925        let json = serde_json::to_string(&config).unwrap();
926        let deserialized: FederationConfig = serde_json::from_str(&json).unwrap();
927        assert_eq!(deserialized.sync_interval_secs, 300);
928        assert!(deserialized.auto_resolve_conflicts);
929    }
930
931    #[test]
932    fn test_peer_config_zeroize_on_drop() {
933        let peer = PeerConfig {
934            node_id: "peer1".to_string(),
935            name: "Test Peer".to_string(),
936            endpoint: "https://example.com".to_string(),
937            api_key: Some("secret-key-12345".to_string()),
938            enabled: true,
939        };
940        // Drop should zeroize the api_key
941        drop(peer);
942    }
943
944    #[test]
945    fn test_peer_config_without_api_key() {
946        let peer = PeerConfig {
947            node_id: "peer1".to_string(),
948            name: "Test Peer".to_string(),
949            endpoint: "https://example.com".to_string(),
950            api_key: None,
951            enabled: false,
952        };
953        assert!(!peer.enabled);
954        drop(peer);
955    }
956
957    #[test]
958    fn test_peer_config_serialization() {
959        let peer = PeerConfig {
960            node_id: "p1".to_string(),
961            name: "Peer 1".to_string(),
962            endpoint: "https://peer1.example.com".to_string(),
963            api_key: Some("key".to_string()),
964            enabled: true,
965        };
966        let json = serde_json::to_string(&peer).unwrap();
967        let deserialized: PeerConfig = serde_json::from_str(&json).unwrap();
968        assert_eq!(deserialized.node_id, "p1");
969        assert_eq!(deserialized.api_key.as_deref(), Some("key"));
970    }
971
972    #[test]
973    fn test_model_sync_state_default_fields() {
974        let state = ModelSyncState {
975            name: "test-model".to_string(),
976            clock: VectorClock::new(),
977            last_sync: HashMap::new(),
978            known_versions: HashSet::new(),
979            pending_upload: HashSet::new(),
980            pending_download: HashSet::new(),
981        };
982        assert_eq!(state.name, "test-model");
983        assert!(state.known_versions.is_empty());
984    }
985
986    #[test]
987    fn test_sync_manifest_serialization() {
988        let manifest = SyncManifest {
989            source_node: "node-a".to_string(),
990            timestamp: Utc::now(),
991            models: vec![],
992            clock: VectorClock::new(),
993        };
994        let json = serde_json::to_string(&manifest).unwrap();
995        let deserialized: SyncManifest = serde_json::from_str(&json).unwrap();
996        assert_eq!(deserialized.source_node, "node-a");
997    }
998
999    #[test]
1000    fn test_sync_result_serialization() {
1001        let result = SyncResult {
1002            peer_id: "peer-1".to_string(),
1003            timestamp: Utc::now(),
1004            duration_ms: 1500,
1005            models_synced: 3,
1006            versions_uploaded: 2,
1007            versions_downloaded: 1,
1008            conflicts: vec![],
1009            errors: vec![],
1010        };
1011        let json = serde_json::to_string(&result).unwrap();
1012        assert!(json.contains("peer-1"));
1013    }
1014
1015    #[test]
1016    fn test_sync_conflict_serialization() {
1017        let conflict = SyncConflict {
1018            model: "llama".to_string(),
1019            local_version: "abc".to_string(),
1020            remote_version: "def".to_string(),
1021            remote_node: "node-b".to_string(),
1022            resolution: Some(ConflictResolution::KeepLocal),
1023        };
1024        let json = serde_json::to_string(&conflict).unwrap();
1025        assert!(json.contains("KeepLocal"));
1026    }
1027
1028    #[test]
1029    fn test_conflict_resolution_variants() {
1030        let keep = ConflictResolution::KeepLocal;
1031        let remote = ConflictResolution::UseRemote;
1032        let branch = ConflictResolution::Branch {
1033            local_name: "local-v2".to_string(),
1034            remote_name: "remote-v2".to_string(),
1035        };
1036        let manual = ConflictResolution::Manual;
1037
1038        // Serialization round-trip
1039        for r in [keep, remote, manual] {
1040            let json = serde_json::to_string(&r).unwrap();
1041            let _: ConflictResolution = serde_json::from_str(&json).unwrap();
1042        }
1043        let json = serde_json::to_string(&branch).unwrap();
1044        assert!(json.contains("local-v2"));
1045    }
1046
1047    #[test]
1048    fn test_version_manifest_entry() {
1049        let entry = VersionManifestEntry {
1050            version: 3,
1051            checkpoint_id: "cp-xyz".to_string(),
1052            created_at: Utc::now(),
1053            checksum: "sha256-abc".to_string(),
1054            size_bytes: 1024,
1055            parent_id: Some("cp-prev".to_string()),
1056            origin_node: "node-a".to_string(),
1057        };
1058        let json = serde_json::to_string(&entry).unwrap();
1059        let deserialized: VersionManifestEntry = serde_json::from_str(&json).unwrap();
1060        assert_eq!(deserialized.version, 3);
1061        assert_eq!(deserialized.checkpoint_id, "cp-xyz");
1062    }
1063
1064    #[test]
1065    fn test_federation_status_serialization() {
1066        let status = FederationStatus {
1067            node_id: "n1".to_string(),
1068            node_name: "Node One".to_string(),
1069            peer_count: 2,
1070            model_count: 5,
1071            last_sync: Some(Utc::now()),
1072            clock: VectorClock::new(),
1073        };
1074        let json = serde_json::to_string(&status).unwrap();
1075        let d: FederationStatus = serde_json::from_str(&json).unwrap();
1076        assert_eq!(d.peer_count, 2);
1077    }
1078
1079    #[tokio::test]
1080    async fn test_federation_manager_new() {
1081        let dir = tempfile::tempdir().unwrap();
1082        let config = FederationConfig::default();
1083        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1084        assert!(!mgr.node_id().is_empty());
1085        assert!(mgr.peers().is_empty());
1086    }
1087
1088    #[tokio::test]
1089    async fn test_federation_manager_add_remove_peer() {
1090        let dir = tempfile::tempdir().unwrap();
1091        let config = FederationConfig::default();
1092        let mut mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1093
1094        let peer = PeerConfig {
1095            node_id: "peer-1".to_string(),
1096            name: "Peer 1".to_string(),
1097            endpoint: "https://peer1.example.com".to_string(),
1098            api_key: None,
1099            enabled: true,
1100        };
1101        mgr.add_peer(peer);
1102        assert_eq!(mgr.peers().len(), 1);
1103
1104        mgr.remove_peer("peer-1");
1105        assert!(mgr.peers().is_empty());
1106    }
1107
1108    #[tokio::test]
1109    async fn test_federation_manager_status() {
1110        let dir = tempfile::tempdir().unwrap();
1111        let config = FederationConfig {
1112            node_name: "test-node".to_string(),
1113            ..Default::default()
1114        };
1115        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1116
1117        let status = mgr.status().await;
1118        assert_eq!(status.node_name, "test-node");
1119        assert_eq!(status.peer_count, 0);
1120        assert_eq!(status.model_count, 0);
1121        assert!(status.last_sync.is_none());
1122    }
1123
1124    #[tokio::test]
1125    async fn test_federation_manager_get_history_empty() {
1126        let dir = tempfile::tempdir().unwrap();
1127        let config = FederationConfig::default();
1128        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1129
1130        let history = mgr.get_history(None).await;
1131        assert!(history.is_empty());
1132
1133        let limited = mgr.get_history(Some(5)).await;
1134        assert!(limited.is_empty());
1135    }
1136
1137    #[tokio::test]
1138    async fn test_federation_manager_generate_manifest_empty() {
1139        let dir = tempfile::tempdir().unwrap();
1140        let config = FederationConfig::default();
1141        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1142
1143        let manifest = mgr.generate_manifest(vec![]).await;
1144        assert!(manifest.models.is_empty());
1145        assert_eq!(manifest.source_node, mgr.node_id());
1146    }
1147
1148    #[tokio::test]
1149    async fn test_federation_manager_generate_manifest_with_models() {
1150        let dir = tempfile::tempdir().unwrap();
1151        let config = FederationConfig::default();
1152        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1153
1154        let versions = vec![crate::version::ModelVersion {
1155            version: 1,
1156            checkpoint_id: "cp-1".to_string(),
1157            timestamp: Utc::now(),
1158            checksum_sha256: "abc123".to_string(),
1159            size_bytes: 1024,
1160            compressed_size_bytes: 800,
1161            parent_version: None,
1162            format: "safetensors".to_string(),
1163            metadata: HashMap::new(),
1164            file_path: "models/cp-1.enc".to_string(),
1165        }];
1166
1167        let manifest = mgr
1168            .generate_manifest(vec![("my-model".to_string(), versions)])
1169            .await;
1170        assert_eq!(manifest.models.len(), 1);
1171        assert_eq!(manifest.models[0].name, "my-model");
1172        assert_eq!(manifest.models[0].versions.len(), 1);
1173    }
1174
1175    #[test]
1176    fn test_compute_delta_empty_manifests() {
1177        let dir = tempfile::tempdir().unwrap();
1178        let config = FederationConfig::default();
1179        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1180
1181        let local = SyncManifest {
1182            source_node: "a".to_string(),
1183            timestamp: Utc::now(),
1184            models: vec![],
1185            clock: VectorClock::new(),
1186        };
1187        let remote = SyncManifest {
1188            source_node: "b".to_string(),
1189            timestamp: Utc::now(),
1190            models: vec![],
1191            clock: VectorClock::new(),
1192        };
1193
1194        let delta = mgr.compute_delta(&local, &remote);
1195        assert!(delta.to_upload.is_empty());
1196        assert!(delta.to_download.is_empty());
1197        assert!(delta.conflicts.is_empty());
1198    }
1199
1200    #[test]
1201    fn test_compute_delta_local_only_model() {
1202        let dir = tempfile::tempdir().unwrap();
1203        let config = FederationConfig::default();
1204        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1205
1206        let local = SyncManifest {
1207            source_node: "a".to_string(),
1208            timestamp: Utc::now(),
1209            models: vec![ModelManifestEntry {
1210                name: "model-x".to_string(),
1211                versions: vec![VersionManifestEntry {
1212                    version: 1,
1213                    checkpoint_id: "cp-1".to_string(),
1214                    created_at: Utc::now(),
1215                    checksum: "sha".to_string(),
1216                    size_bytes: 100,
1217                    parent_id: None,
1218                    origin_node: "a".to_string(),
1219                }],
1220                clock: VectorClock::new(),
1221            }],
1222            clock: VectorClock::new(),
1223        };
1224        let remote = SyncManifest {
1225            source_node: "b".to_string(),
1226            timestamp: Utc::now(),
1227            models: vec![],
1228            clock: VectorClock::new(),
1229        };
1230
1231        let delta = mgr.compute_delta(&local, &remote);
1232        assert_eq!(delta.to_upload.len(), 1);
1233        assert_eq!(delta.to_upload[0].model, "model-x");
1234        assert!(delta.to_download.is_empty());
1235    }
1236
1237    #[test]
1238    fn test_compute_delta_remote_only_model() {
1239        let dir = tempfile::tempdir().unwrap();
1240        let config = FederationConfig::default();
1241        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1242
1243        let local = SyncManifest {
1244            source_node: "a".to_string(),
1245            timestamp: Utc::now(),
1246            models: vec![],
1247            clock: VectorClock::new(),
1248        };
1249        let remote = SyncManifest {
1250            source_node: "b".to_string(),
1251            timestamp: Utc::now(),
1252            models: vec![ModelManifestEntry {
1253                name: "model-y".to_string(),
1254                versions: vec![VersionManifestEntry {
1255                    version: 1,
1256                    checkpoint_id: "cp-2".to_string(),
1257                    created_at: Utc::now(),
1258                    checksum: "sha2".to_string(),
1259                    size_bytes: 200,
1260                    parent_id: None,
1261                    origin_node: "b".to_string(),
1262                }],
1263                clock: VectorClock::new(),
1264            }],
1265            clock: VectorClock::new(),
1266        };
1267
1268        let delta = mgr.compute_delta(&local, &remote);
1269        assert!(delta.to_upload.is_empty());
1270        assert_eq!(delta.to_download.len(), 1);
1271        assert_eq!(delta.to_download[0].model, "model-y");
1272    }
1273
1274    #[test]
1275    fn test_compute_delta_conflict() {
1276        let dir = tempfile::tempdir().unwrap();
1277        let config = FederationConfig::default();
1278        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1279
1280        let local = SyncManifest {
1281            source_node: "a".to_string(),
1282            timestamp: Utc::now(),
1283            models: vec![ModelManifestEntry {
1284                name: "model-z".to_string(),
1285                versions: vec![VersionManifestEntry {
1286                    version: 1,
1287                    checkpoint_id: "cp-local".to_string(),
1288                    created_at: Utc::now(),
1289                    checksum: "sha-l".to_string(),
1290                    size_bytes: 100,
1291                    parent_id: None,
1292                    origin_node: "a".to_string(),
1293                }],
1294                clock: VectorClock::new(),
1295            }],
1296            clock: VectorClock::new(),
1297        };
1298        let remote = SyncManifest {
1299            source_node: "b".to_string(),
1300            timestamp: Utc::now(),
1301            models: vec![ModelManifestEntry {
1302                name: "model-z".to_string(),
1303                versions: vec![VersionManifestEntry {
1304                    version: 1,
1305                    checkpoint_id: "cp-remote".to_string(),
1306                    created_at: Utc::now(),
1307                    checksum: "sha-r".to_string(),
1308                    size_bytes: 100,
1309                    parent_id: None,
1310                    origin_node: "b".to_string(),
1311                }],
1312                clock: VectorClock::new(),
1313            }],
1314            clock: VectorClock::new(),
1315        };
1316
1317        let delta = mgr.compute_delta(&local, &remote);
1318        assert_eq!(delta.conflicts.len(), 1);
1319        assert_eq!(delta.conflicts[0].model, "model-z");
1320    }
1321
1322    #[test]
1323    fn test_compute_delta_shared_versions() {
1324        let dir = tempfile::tempdir().unwrap();
1325        let config = FederationConfig::default();
1326        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1327
1328        // Both nodes have the same version — no sync needed
1329        let entry = VersionManifestEntry {
1330            version: 1,
1331            checkpoint_id: "cp-same".to_string(),
1332            created_at: Utc::now(),
1333            checksum: "sha".to_string(),
1334            size_bytes: 100,
1335            parent_id: None,
1336            origin_node: "a".to_string(),
1337        };
1338        let local = SyncManifest {
1339            source_node: "a".to_string(),
1340            timestamp: Utc::now(),
1341            models: vec![ModelManifestEntry {
1342                name: "m".to_string(),
1343                versions: vec![entry.clone()],
1344                clock: VectorClock::new(),
1345            }],
1346            clock: VectorClock::new(),
1347        };
1348        let remote = SyncManifest {
1349            source_node: "b".to_string(),
1350            timestamp: Utc::now(),
1351            models: vec![ModelManifestEntry {
1352                name: "m".to_string(),
1353                versions: vec![entry],
1354                clock: VectorClock::new(),
1355            }],
1356            clock: VectorClock::new(),
1357        };
1358
1359        let delta = mgr.compute_delta(&local, &remote);
1360        assert!(delta.to_upload.is_empty());
1361        assert!(delta.to_download.is_empty());
1362        assert!(delta.conflicts.is_empty());
1363    }
1364
1365    #[tokio::test]
1366    async fn test_federation_manager_load_existing_state() {
1367        let dir = tempfile::tempdir().unwrap();
1368        let state_file = dir.path().join("federation_state.json");
1369        let saved = SavedFederationState {
1370            models: HashMap::new(),
1371            clock: {
1372                let mut c = VectorClock::new();
1373                c.increment("existing");
1374                c
1375            },
1376            history: vec![],
1377        };
1378        std::fs::write(&state_file, serde_json::to_string(&saved).unwrap()).unwrap();
1379
1380        let config = FederationConfig::default();
1381        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1382
1383        let status = mgr.status().await;
1384        assert_eq!(status.clock.timestamps.get("existing"), Some(&1));
1385    }
1386
1387    #[test]
1388    fn test_sync_item_debug() {
1389        let item = SyncItem {
1390            model: "m".to_string(),
1391            checkpoint_id: "cp".to_string(),
1392            size_bytes: 42,
1393        };
1394        let dbg = format!("{:?}", item);
1395        assert!(dbg.contains("42"));
1396    }
1397
1398    #[test]
1399    fn test_model_manifest_entry_serialization() {
1400        let entry = ModelManifestEntry {
1401            name: "llm".to_string(),
1402            versions: vec![],
1403            clock: VectorClock::new(),
1404        };
1405        let json = serde_json::to_string(&entry).unwrap();
1406        let d: ModelManifestEntry = serde_json::from_str(&json).unwrap();
1407        assert_eq!(d.name, "llm");
1408    }
1409
1410    #[tokio::test]
1411    async fn test_federation_manager_save_and_reload_state() {
1412        let dir = tempfile::tempdir().unwrap();
1413        let config = FederationConfig::default();
1414        let node_id = config.node_id.clone();
1415        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1416
1417        // Modify internal state via generate_manifest (increments clock indirectly)
1418        // Then save state by accessing internals
1419        {
1420            let mut state = mgr.state.write().await;
1421            state.clock.increment(&node_id);
1422            state.models.insert(
1423                "saved-model".to_string(),
1424                ModelSyncState {
1425                    name: "saved-model".to_string(),
1426                    clock: VectorClock::new(),
1427                    last_sync: HashMap::new(),
1428                    known_versions: HashSet::new(),
1429                    pending_upload: HashSet::new(),
1430                    pending_download: HashSet::new(),
1431                },
1432            );
1433            mgr.save_state(&state).unwrap();
1434        }
1435
1436        // Reload from same directory
1437        let config2 = FederationConfig::default();
1438        let mgr2 = FederationManager::new(config2, dir.path().to_path_buf()).unwrap();
1439        let status2 = mgr2.status().await;
1440        assert_eq!(status2.model_count, 1);
1441        assert_eq!(status2.clock.timestamps.get(&node_id), Some(&1));
1442    }
1443
1444    #[test]
1445    fn test_compute_delta_mixed_versions() {
1446        // Both nodes have the model but with different extra versions
1447        let dir = tempfile::tempdir().unwrap();
1448        let config = FederationConfig::default();
1449        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1450
1451        let shared = VersionManifestEntry {
1452            version: 1,
1453            checkpoint_id: "cp-shared".to_string(),
1454            created_at: Utc::now(),
1455            checksum: "sha".to_string(),
1456            size_bytes: 100,
1457            parent_id: None,
1458            origin_node: "a".to_string(),
1459        };
1460        let local_only = VersionManifestEntry {
1461            version: 2,
1462            checkpoint_id: "cp-local-v2".to_string(),
1463            created_at: Utc::now(),
1464            checksum: "sha2".to_string(),
1465            size_bytes: 200,
1466            parent_id: Some("cp-shared".to_string()),
1467            origin_node: "a".to_string(),
1468        };
1469        let remote_only = VersionManifestEntry {
1470            version: 3,
1471            checkpoint_id: "cp-remote-v3".to_string(),
1472            created_at: Utc::now(),
1473            checksum: "sha3".to_string(),
1474            size_bytes: 300,
1475            parent_id: Some("cp-shared".to_string()),
1476            origin_node: "b".to_string(),
1477        };
1478
1479        let local = SyncManifest {
1480            source_node: "a".to_string(),
1481            timestamp: Utc::now(),
1482            models: vec![ModelManifestEntry {
1483                name: "m".to_string(),
1484                versions: vec![shared.clone(), local_only],
1485                clock: VectorClock::new(),
1486            }],
1487            clock: VectorClock::new(),
1488        };
1489        let remote = SyncManifest {
1490            source_node: "b".to_string(),
1491            timestamp: Utc::now(),
1492            models: vec![ModelManifestEntry {
1493                name: "m".to_string(),
1494                versions: vec![shared, remote_only],
1495                clock: VectorClock::new(),
1496            }],
1497            clock: VectorClock::new(),
1498        };
1499
1500        let delta = mgr.compute_delta(&local, &remote);
1501        assert_eq!(delta.to_upload.len(), 1);
1502        assert_eq!(delta.to_upload[0].checkpoint_id, "cp-local-v2");
1503        assert_eq!(delta.to_download.len(), 1);
1504        assert_eq!(delta.to_download[0].checkpoint_id, "cp-remote-v3");
1505        assert!(delta.conflicts.is_empty()); // Different version numbers, no conflict
1506    }
1507
1508    #[tokio::test]
1509    async fn test_federation_manager_generate_manifest_multi_version() {
1510        let dir = tempfile::tempdir().unwrap();
1511        let config = FederationConfig::default();
1512        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1513
1514        let versions = vec![
1515            crate::version::ModelVersion {
1516                version: 1,
1517                checkpoint_id: "cp-1".to_string(),
1518                timestamp: Utc::now(),
1519                checksum_sha256: "abc".to_string(),
1520                size_bytes: 512,
1521                compressed_size_bytes: 400,
1522                parent_version: None,
1523                format: "safetensors".to_string(),
1524                metadata: HashMap::new(),
1525                file_path: "models/cp-1.enc".to_string(),
1526            },
1527            crate::version::ModelVersion {
1528                version: 2,
1529                checkpoint_id: "cp-2".to_string(),
1530                timestamp: Utc::now(),
1531                checksum_sha256: "def".to_string(),
1532                size_bytes: 1024,
1533                compressed_size_bytes: 900,
1534                parent_version: Some(1),
1535                format: "gguf".to_string(),
1536                metadata: HashMap::new(),
1537                file_path: "models/cp-2.enc".to_string(),
1538            },
1539        ];
1540
1541        let manifest = mgr
1542            .generate_manifest(vec![("multi-v".to_string(), versions)])
1543            .await;
1544        assert_eq!(manifest.models.len(), 1);
1545        assert_eq!(manifest.models[0].versions.len(), 2);
1546        assert_eq!(manifest.models[0].versions[0].version, 1);
1547        assert_eq!(manifest.models[0].versions[1].version, 2);
1548        assert_eq!(
1549            manifest.models[0].versions[1].parent_id.as_deref(),
1550            Some("1")
1551        );
1552    }
1553
1554    #[test]
1555    fn test_saved_federation_state_roundtrip() {
1556        let mut clock = VectorClock::new();
1557        clock.increment("n1");
1558        clock.increment("n2");
1559
1560        let saved = SavedFederationState {
1561            models: {
1562                let mut m = HashMap::new();
1563                m.insert(
1564                    "model-a".to_string(),
1565                    ModelSyncState {
1566                        name: "model-a".to_string(),
1567                        clock: clock.clone(),
1568                        last_sync: HashMap::new(),
1569                        known_versions: {
1570                            let mut s = HashSet::new();
1571                            s.insert("cp-1".to_string());
1572                            s
1573                        },
1574                        pending_upload: HashSet::new(),
1575                        pending_download: HashSet::new(),
1576                    },
1577                );
1578                m
1579            },
1580            clock,
1581            history: vec![SyncResult {
1582                peer_id: "p1".to_string(),
1583                timestamp: Utc::now(),
1584                duration_ms: 500,
1585                models_synced: 1,
1586                versions_uploaded: 0,
1587                versions_downloaded: 1,
1588                conflicts: vec![],
1589                errors: vec![],
1590            }],
1591        };
1592
1593        let json = serde_json::to_string(&saved).unwrap();
1594        let loaded: SavedFederationState = serde_json::from_str(&json).unwrap();
1595        assert_eq!(loaded.models.len(), 1);
1596        assert!(loaded.models.contains_key("model-a"));
1597        assert_eq!(loaded.history.len(), 1);
1598        assert_eq!(loaded.clock.timestamps.get("n1"), Some(&1));
1599    }
1600
1601    #[test]
1602    fn test_sync_delta_debug() {
1603        let delta = SyncDelta {
1604            to_upload: vec![],
1605            to_download: vec![],
1606            conflicts: vec![],
1607        };
1608        let dbg = format!("{:?}", delta);
1609        assert!(dbg.contains("SyncDelta"));
1610    }
1611
1612    #[test]
1613    fn test_model_sync_state_with_pending() {
1614        let mut state = ModelSyncState {
1615            name: "pending-model".to_string(),
1616            clock: VectorClock::new(),
1617            last_sync: HashMap::new(),
1618            known_versions: HashSet::new(),
1619            pending_upload: HashSet::new(),
1620            pending_download: HashSet::new(),
1621        };
1622        state.pending_upload.insert("cp-up-1".to_string());
1623        state
1624            .pending_download
1625            .insert(("node-b".to_string(), "cp-down-1".to_string()));
1626        state.known_versions.insert("cp-1".to_string());
1627        state.last_sync.insert("node-b".to_string(), Utc::now());
1628
1629        let json = serde_json::to_string(&state).unwrap();
1630        let d: ModelSyncState = serde_json::from_str(&json).unwrap();
1631        assert_eq!(d.pending_upload.len(), 1);
1632        assert_eq!(d.pending_download.len(), 1);
1633        assert_eq!(d.known_versions.len(), 1);
1634        assert_eq!(d.last_sync.len(), 1);
1635    }
1636
1637    #[tokio::test]
1638    async fn test_federation_manager_history_with_limit() {
1639        let dir = tempfile::tempdir().unwrap();
1640        let state_file = dir.path().join("federation_state.json");
1641
1642        // Pre-populate with history
1643        let history: Vec<SyncResult> = (0..5)
1644            .map(|i| SyncResult {
1645                peer_id: format!("peer-{}", i),
1646                timestamp: Utc::now(),
1647                duration_ms: 100 * (i as u64 + 1),
1648                models_synced: 1,
1649                versions_uploaded: 0,
1650                versions_downloaded: 1,
1651                conflicts: vec![],
1652                errors: vec![],
1653            })
1654            .collect();
1655
1656        let saved = SavedFederationState {
1657            models: HashMap::new(),
1658            clock: VectorClock::new(),
1659            history,
1660        };
1661        std::fs::write(&state_file, serde_json::to_string(&saved).unwrap()).unwrap();
1662
1663        let config = FederationConfig::default();
1664        let mgr = FederationManager::new(config, dir.path().to_path_buf()).unwrap();
1665
1666        let all = mgr.get_history(None).await;
1667        assert_eq!(all.len(), 5);
1668
1669        let limited = mgr.get_history(Some(3)).await;
1670        assert_eq!(limited.len(), 3);
1671        // get_history returns reversed (most recent first)
1672        assert_eq!(limited[0].peer_id, "peer-4");
1673    }
1674
1675    #[test]
1676    fn test_conflict_resolution_branch_serialization() {
1677        let res = ConflictResolution::Branch {
1678            local_name: "main-v2".to_string(),
1679            remote_name: "fork-v2".to_string(),
1680        };
1681        let json = serde_json::to_string(&res).unwrap();
1682        let d: ConflictResolution = serde_json::from_str(&json).unwrap();
1683        match d {
1684            ConflictResolution::Branch {
1685                local_name,
1686                remote_name,
1687            } => {
1688                assert_eq!(local_name, "main-v2");
1689                assert_eq!(remote_name, "fork-v2");
1690            }
1691            _ => panic!("Expected Branch variant"),
1692        }
1693    }
1694
1695    #[test]
1696    fn test_federation_status_without_last_sync() {
1697        let status = FederationStatus {
1698            node_id: "n1".to_string(),
1699            node_name: "node1".to_string(),
1700            peer_count: 0,
1701            model_count: 0,
1702            last_sync: None,
1703            clock: VectorClock::new(),
1704        };
1705        let json = serde_json::to_string(&status).unwrap();
1706        let d: FederationStatus = serde_json::from_str(&json).unwrap();
1707        assert!(d.last_sync.is_none());
1708    }
1709}
ai_model_vault/federation.rs

ai_model_vault/
federation.rs
