ultrafast_mcp_core/utils/
progress.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::time::{Duration, SystemTime, UNIX_EPOCH};
4
5/// Progress information for long-running operations
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct Progress {
8    /// Unique identifier for this progress
9    pub id: String,
10    /// Current progress value
11    pub current: u64,
12    /// Total expected value (if known)
13    #[serde(skip_serializing_if = "Option::is_none")]
14    pub total: Option<u64>,
15    /// Human-readable description
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub description: Option<String>,
18    /// Current status
19    pub status: ProgressStatus,
20    /// Timestamp of last update (Unix timestamp)
21    pub updated_at: u64,
22    /// Additional metadata
23    #[serde(skip_serializing_if = "HashMap::is_empty", default)]
24    pub metadata: HashMap<String, serde_json::Value>,
25}
26
27/// Progress status
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29#[serde(rename_all = "lowercase")]
30pub enum ProgressStatus {
31    /// Operation is starting
32    Starting,
33    /// Operation is in progress
34    Running,
35    /// Operation completed successfully
36    Completed,
37    /// Operation failed
38    Failed,
39    /// Operation was cancelled
40    Cancelled,
41    /// Operation is paused
42    Paused,
43}
44
45impl Progress {
46    /// Create a new progress instance
47    pub fn new(id: impl Into<String>) -> Self {
48        Self {
49            id: id.into(),
50            current: 0,
51            total: None,
52            description: None,
53            status: ProgressStatus::Starting,
54            updated_at: current_timestamp(),
55            metadata: HashMap::new(),
56        }
57    }
58
59    /// Set the total expected value
60    pub fn with_total(mut self, total: u64) -> Self {
61        self.total = Some(total);
62        self.updated_at = current_timestamp();
63        self
64    }
65
66    /// Set the description
67    pub fn with_description(mut self, description: impl Into<String>) -> Self {
68        self.description = Some(description.into());
69        self.updated_at = current_timestamp();
70        self
71    }
72
73    /// Set the current progress
74    pub fn with_current(mut self, current: u64) -> Self {
75        self.current = current;
76        self.updated_at = current_timestamp();
77        self
78    }
79
80    /// Set the status
81    pub fn with_status(mut self, status: ProgressStatus) -> Self {
82        self.status = status;
83        self.updated_at = current_timestamp();
84        self
85    }
86
87    /// Add metadata
88    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
89        self.metadata.insert(key.into(), value);
90        self.updated_at = current_timestamp();
91        self
92    }
93
94    /// Update the current progress
95    pub fn update(&mut self, current: u64) {
96        self.current = current;
97        self.updated_at = current_timestamp();
98    }
99
100    /// Update with description
101    pub fn update_with_description(&mut self, current: u64, description: impl Into<String>) {
102        self.current = current;
103        self.description = Some(description.into());
104        self.updated_at = current_timestamp();
105    }
106
107    /// Mark as completed
108    pub fn complete(&mut self) {
109        if let Some(total) = self.total {
110            self.current = total;
111        }
112        self.status = ProgressStatus::Completed;
113        self.updated_at = current_timestamp();
114    }
115
116    /// Mark as failed
117    pub fn fail(&mut self, error: Option<String>) {
118        self.status = ProgressStatus::Failed;
119        if let Some(error) = error {
120            self.metadata
121                .insert("error".to_string(), serde_json::Value::String(error));
122        }
123        self.updated_at = current_timestamp();
124    }
125
126    /// Mark as cancelled
127    pub fn cancel(&mut self) {
128        self.status = ProgressStatus::Cancelled;
129        self.updated_at = current_timestamp();
130    }
131
132    /// Get progress percentage (0-100) if total is known
133    pub fn percentage(&self) -> Option<f64> {
134        self.total.map(|total| {
135            if total == 0 {
136                100.0
137            } else {
138                (self.current as f64 / total as f64) * 100.0
139            }
140        })
141    }
142
143    /// Check if the operation is finished
144    pub fn is_finished(&self) -> bool {
145        matches!(
146            self.status,
147            ProgressStatus::Completed | ProgressStatus::Failed | ProgressStatus::Cancelled
148        )
149    }
150
151    /// Check if the operation is active
152    pub fn is_active(&self) -> bool {
153        matches!(
154            self.status,
155            ProgressStatus::Running | ProgressStatus::Starting
156        )
157    }
158
159    /// Get the age of this progress update
160    pub fn age(&self) -> Duration {
161        let now = current_timestamp();
162        Duration::from_secs(now.saturating_sub(self.updated_at))
163    }
164}
165
166/// Progress tracker for managing multiple progress instances
167#[derive(Debug, Default)]
168pub struct ProgressTracker {
169    progress_map: HashMap<String, Progress>,
170}
171
172impl ProgressTracker {
173    /// Create a new progress tracker
174    pub fn new() -> Self {
175        Self::default()
176    }
177
178    /// Start tracking a new progress
179    pub fn start(&mut self, id: impl Into<String>) -> &mut Progress {
180        let id = id.into();
181        let progress = Progress::new(id.clone()).with_status(ProgressStatus::Running);
182        self.progress_map.entry(id.clone()).or_insert(progress)
183    }
184
185    /// Get a progress by ID
186    pub fn get(&self, id: &str) -> Option<&Progress> {
187        self.progress_map.get(id)
188    }
189
190    /// Get a mutable progress by ID
191    pub fn get_mut(&mut self, id: &str) -> Option<&mut Progress> {
192        self.progress_map.get_mut(id)
193    }
194
195    /// Update progress
196    pub fn update(&mut self, id: &str, current: u64) -> Option<&Progress> {
197        let progress = self.progress_map.get_mut(id)?;
198        progress.update(current);
199        Some(progress)
200    }
201
202    /// Complete a progress
203    pub fn complete(&mut self, id: &str) -> Option<&Progress> {
204        let progress = self.progress_map.get_mut(id)?;
205        progress.complete();
206        Some(progress)
207    }
208
209    /// Fail a progress
210    pub fn fail(&mut self, id: &str, error: Option<String>) -> Option<&Progress> {
211        let progress = self.progress_map.get_mut(id)?;
212        progress.fail(error);
213        Some(progress)
214    }
215
216    /// Cancel a progress
217    pub fn cancel(&mut self, id: &str) -> Option<&Progress> {
218        let progress = self.progress_map.get_mut(id)?;
219        progress.cancel();
220        Some(progress)
221    }
222
223    /// Remove a progress
224    pub fn remove(&mut self, id: &str) -> Option<Progress> {
225        self.progress_map.remove(id)
226    }
227
228    /// Get all progress instances
229    pub fn all(&self) -> impl Iterator<Item = &Progress> {
230        self.progress_map.values()
231    }
232
233    /// Get all active progress instances
234    pub fn active(&self) -> impl Iterator<Item = &Progress> {
235        self.progress_map.values().filter(|p| p.is_active())
236    }
237
238    /// Get all finished progress instances
239    pub fn finished(&self) -> impl Iterator<Item = &Progress> {
240        self.progress_map.values().filter(|p| p.is_finished())
241    }
242
243    /// Clean up old finished progress instances
244    pub fn cleanup_finished(&mut self, max_age: Duration) {
245        let cutoff = current_timestamp() - max_age.as_secs();
246        self.progress_map
247            .retain(|_, progress| !progress.is_finished() || progress.updated_at > cutoff);
248    }
249
250    /// Clean up all progress entries older than max_age (including active ones)
251    pub fn cleanup_all_old(&mut self, max_age: Duration) {
252        let cutoff = current_timestamp() - max_age.as_secs();
253        self.progress_map
254            .retain(|_, progress| progress.updated_at > cutoff);
255    }
256
257    /// Clean up progress entries that have been inactive for too long
258    pub fn cleanup_inactive(&mut self, max_inactive_age: Duration) {
259        let cutoff = current_timestamp() - max_inactive_age.as_secs();
260        self.progress_map
261            .retain(|_, progress| progress.is_active() || progress.updated_at > cutoff);
262    }
263
264    /// Get the number of tracked progress instances
265    pub fn len(&self) -> usize {
266        self.progress_map.len()
267    }
268
269    /// Check if the tracker is empty
270    pub fn is_empty(&self) -> bool {
271        self.progress_map.is_empty()
272    }
273}
274
275/// Get current Unix timestamp
276fn current_timestamp() -> u64 {
277    SystemTime::now()
278        .duration_since(UNIX_EPOCH)
279        .unwrap_or_default()
280        .as_secs()
281}
282
283/// Progress notification for MCP protocol
284#[derive(Debug, Clone, Serialize, Deserialize)]
285pub struct ProgressNotification {
286    /// Progress information
287    #[serde(flatten)]
288    pub progress: Progress,
289}
290
291impl ProgressNotification {
292    /// Create a new progress notification
293    pub fn new(progress: Progress) -> Self {
294        Self { progress }
295    }
296}
297
298impl From<Progress> for ProgressNotification {
299    fn from(progress: Progress) -> Self {
300        Self::new(progress)
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307
308    #[test]
309    fn test_progress_creation() {
310        let progress = Progress::new("test")
311            .with_total(100)
312            .with_description("Test operation")
313            .with_current(25);
314
315        assert_eq!(progress.id, "test");
316        assert_eq!(progress.current, 25);
317        assert_eq!(progress.total, Some(100));
318        assert_eq!(progress.description, Some("Test operation".to_string()));
319        assert_eq!(progress.status, ProgressStatus::Starting);
320    }
321
322    #[test]
323    fn test_progress_percentage() {
324        let mut progress = Progress::new("test").with_total(100);
325        progress.update(25);
326        assert_eq!(progress.percentage(), Some(25.0));
327
328        progress.update(50);
329        assert_eq!(progress.percentage(), Some(50.0));
330
331        let progress_no_total = Progress::new("test");
332        assert_eq!(progress_no_total.percentage(), None);
333    }
334
335    #[test]
336    fn test_progress_status() {
337        let mut progress = Progress::new("test").with_status(ProgressStatus::Running);
338        assert!(progress.is_active());
339        assert!(!progress.is_finished());
340
341        progress.complete();
342        assert!(progress.is_finished());
343        assert!(!progress.is_active());
344        assert_eq!(progress.status, ProgressStatus::Completed);
345    }
346
347    #[test]
348    fn test_progress_tracker() {
349        let mut tracker = ProgressTracker::new();
350
351        // Start tracking
352        let progress = tracker.start("test1");
353        progress.update(50);
354
355        // Check active
356        assert_eq!(tracker.active().count(), 1);
357        assert_eq!(tracker.finished().count(), 0);
358
359        // Complete
360        tracker.complete("test1");
361        assert_eq!(tracker.active().count(), 0);
362        assert_eq!(tracker.finished().count(), 1);
363
364        // Test update
365        tracker.start("test2");
366        let updated = tracker.update("test2", 75);
367        assert!(updated.is_some());
368        assert_eq!(updated.unwrap().current, 75);
369    }
370
371    #[test]
372    fn test_progress_cleanup() {
373        let mut tracker = ProgressTracker::new();
374
375        // Add some progress
376        tracker.start("test1");
377        tracker.complete("test1");
378
379        // Cleanup should remove finished progress
380        tracker.cleanup_finished(Duration::from_secs(0));
381        assert_eq!(tracker.len(), 0);
382    }
383}
ultrafast_mcp_core/utils/progress.rs

ultrafast_mcp_core/utils/
progress.rs
