offline_intelligence/engine_management/
download_progress.rs

1//! Engine Download Progress Tracking
2//!
3//! Provides real-time progress tracking specifically for engine downloads.
4//! This is separate from model download tracking to maintain proper type safety
5//! and field naming consistency.
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::sync::Arc;
10use tokio::sync::RwLock;
11use tracing::{debug, info};
12use uuid::Uuid;
13
14/// Status of an engine download
15#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
16pub enum EngineDownloadStatus {
17    Queued,
18    Starting,
19    Downloading,
20    Paused,
21    Extracting,
22    Verifying,
23    Completed,
24    Failed,
25    Cancelled,
26}
27
28impl EngineDownloadStatus {
29    pub fn is_active(&self) -> bool {
30        matches!(self, 
31            EngineDownloadStatus::Queued | 
32            EngineDownloadStatus::Starting | 
33            EngineDownloadStatus::Downloading |
34            EngineDownloadStatus::Paused |
35            EngineDownloadStatus::Extracting |
36            EngineDownloadStatus::Verifying
37        )
38    }
39
40    pub fn is_finished(&self) -> bool {
41        matches!(self, 
42            EngineDownloadStatus::Completed | 
43            EngineDownloadStatus::Failed | 
44            EngineDownloadStatus::Cancelled
45        )
46    }
47}
48
49/// Engine download progress information
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct EngineDownloadProgress {
52    pub download_id: String,
53    pub engine_id: String,
54    pub engine_name: String,
55    pub status: EngineDownloadStatus,
56    pub bytes_downloaded: u64,
57    pub total_bytes: u64,
58    pub progress_percentage: f32,
59    pub speed_bps: f64,
60    pub error_message: Option<String>,
61    pub started_at: Option<chrono::DateTime<chrono::Utc>>,
62    pub completed_at: Option<chrono::DateTime<chrono::Utc>>,
63}
64
65impl Default for EngineDownloadProgress {
66    fn default() -> Self {
67        Self {
68            download_id: String::new(),
69            engine_id: String::new(),
70            engine_name: String::new(),
71            status: EngineDownloadStatus::Queued,
72            bytes_downloaded: 0,
73            total_bytes: 0,
74            progress_percentage: 0.0,
75            speed_bps: 0.0,
76            error_message: None,
77            started_at: None,
78            completed_at: None,
79        }
80    }
81}
82
83/// Progress tracker specifically for engine downloads
84pub struct EngineDownloadProgressTracker {
85    downloads: Arc<RwLock<HashMap<String, EngineDownloadProgress>>>,
86}
87
88impl EngineDownloadProgressTracker {
89    pub fn new() -> Self {
90        Self {
91            downloads: Arc::new(RwLock::new(HashMap::new())),
92        }
93    }
94
95    /// Start tracking a new engine download
96    pub async fn start_download(
97        &self,
98        engine_id: String,
99        engine_name: String,
100        total_bytes: u64,
101    ) -> String {
102        let download_id = Uuid::new_v4().to_string();
103        let progress = EngineDownloadProgress {
104            download_id: download_id.clone(),
105            engine_id: engine_id.clone(),
106            engine_name,
107            status: EngineDownloadStatus::Starting,
108            bytes_downloaded: 0,
109            total_bytes,
110            progress_percentage: 0.0,
111            speed_bps: 0.0,
112            error_message: None,
113            started_at: Some(chrono::Utc::now()),
114            completed_at: None,
115        };
116
117        {
118            let mut downloads = self.downloads.write().await;
119            downloads.insert(engine_id.clone(), progress);
120        }
121
122        info!("Started tracking engine download: {} for engine: {}", download_id, engine_id);
123        download_id
124    }
125
126    /// Update engine download progress
127    pub async fn update_progress(
128        &self,
129        engine_id: &str,
130        bytes_downloaded: u64,
131        status: EngineDownloadStatus,
132        error_message: Option<String>,
133    ) {
134        let mut downloads = self.downloads.write().await;
135        
136        if let Some(progress) = downloads.get_mut(engine_id) {
137            progress.bytes_downloaded = bytes_downloaded;
138            progress.status = status;
139            progress.error_message = error_message;
140
141            // Calculate percentage
142            if progress.total_bytes > 0 {
143                progress.progress_percentage = (bytes_downloaded as f32 / progress.total_bytes as f32) * 100.0;
144            }
145
146            // Update completed timestamp if finished
147            if progress.status.is_finished() {
148                progress.completed_at = Some(chrono::Utc::now());
149            }
150
151            debug!("Updated engine download progress for {}: {:.1}%", engine_id, progress.progress_percentage);
152        }
153    }
154
155    /// Update download status without changing bytes
156    pub async fn update_status(
157        &self,
158        engine_id: &str,
159        status: EngineDownloadStatus,
160    ) {
161        let mut downloads = self.downloads.write().await;
162        
163        if let Some(progress) = downloads.get_mut(engine_id) {
164            progress.status = status;
165            
166            if progress.status.is_finished() {
167                progress.completed_at = Some(chrono::Utc::now());
168            }
169        }
170    }
171
172    /// Get current progress for an engine download
173    pub async fn get_progress(&self, engine_id: &str) -> Option<EngineDownloadProgress> {
174        let downloads = self.downloads.read().await;
175        downloads.get(engine_id).cloned()
176    }
177
178    /// Get all active engine downloads
179    pub async fn get_active_downloads(&self) -> Vec<EngineDownloadProgress> {
180        let downloads = self.downloads.read().await;
181        downloads.values()
182            .filter(|p| p.status.is_active())
183            .cloned()
184            .collect()
185    }
186
187    /// Get all engine downloads (active and completed)
188    pub async fn get_all_downloads(&self) -> Vec<EngineDownloadProgress> {
189        let downloads = self.downloads.read().await;
190        downloads.values().cloned().collect()
191    }
192
193    /// Remove a completed/failed download from tracking
194    pub async fn remove_download(&self, engine_id: &str) {
195        let mut downloads = self.downloads.write().await;
196        if let Some(progress) = downloads.get(engine_id) {
197            if progress.status.is_finished() {
198                downloads.remove(engine_id);
199                info!("Removed completed engine download tracking: {}", engine_id);
200            }
201        }
202    }
203
204    /// Cancel an active download
205    pub async fn cancel_download(&self, engine_id: &str) -> bool {
206        let mut downloads = self.downloads.write().await;
207        if let Some(progress) = downloads.get_mut(engine_id) {
208            if progress.status.is_active() {
209                progress.status = EngineDownloadStatus::Cancelled;
210                progress.completed_at = Some(chrono::Utc::now());
211                info!("Cancelled engine download: {}", engine_id);
212                true
213            } else {
214                false
215            }
216        } else {
217            false
218        }
219    }
220
221    /// Clean up old completed downloads (call periodically)
222    pub async fn cleanup_old_downloads(&self, max_age_hours: i64) {
223        let cutoff = chrono::Utc::now() - chrono::Duration::hours(max_age_hours);
224        let mut downloads = self.downloads.write().await;
225        
226        let to_remove: Vec<String> = downloads.iter()
227            .filter(|(_, progress)| {
228                progress.status.is_finished() && 
229                progress.completed_at.map_or(false, |t| t < cutoff)
230            })
231            .map(|(id, _)| id.clone())
232            .collect();
233        
234        for id in to_remove {
235            downloads.remove(&id);
236        }
237    }
238}
239
240impl Default for EngineDownloadProgressTracker {
241    fn default() -> Self {
242        Self::new()
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    #[tokio::test]
251    async fn test_engine_download_progress_tracking() {
252        let tracker = EngineDownloadProgressTracker::new();
253        
254        let download_id = tracker.start_download(
255            "test-engine".to_string(),
256            "Test Engine".to_string(),
257            1000
258        ).await;
259
260        assert!(!download_id.is_empty());
261
262        // Update progress
263        tracker.update_progress("test-engine", 500, EngineDownloadStatus::Downloading, None).await;
264        
265        let progress = tracker.get_progress("test-engine").await.unwrap();
266        assert_eq!(progress.bytes_downloaded, 500);
267        assert_eq!(progress.progress_percentage, 50.0);
268        assert_eq!(progress.status, EngineDownloadStatus::Downloading);
269        assert_eq!(progress.engine_id, "test-engine");
270    }
271
272    #[tokio::test]
273    async fn test_active_vs_finished() {
274        let tracker = EngineDownloadProgressTracker::new();
275        
276        tracker.start_download(
277            "active-engine".to_string(),
278            "Active Engine".to_string(),
279            1000
280        ).await;
281
282        tracker.update_progress("active-engine", 500, EngineDownloadStatus::Downloading, None).await;
283        
284        let active = tracker.get_active_downloads().await;
285        assert_eq!(active.len(), 1);
286
287        // Complete the download
288        tracker.update_progress("active-engine", 1000, EngineDownloadStatus::Completed, None).await;
289        
290        let active = tracker.get_active_downloads().await;
291        assert_eq!(active.len(), 0);
292        
293        let all = tracker.get_all_downloads().await;
294        assert_eq!(all.len(), 1);
295    }
296}
offline_intelligence/engine_management/download_progress.rs

offline_intelligence/engine_management/
download_progress.rs
