Skip to main content

mcp_kit/server/
progress.rs

1//! Progress tracking for long-running operations.
2//!
3//! When a client sends a request with a `_meta.progressToken`, the server can
4//! send progress notifications back to the client during the operation.
5//!
6//! # Example
7//! ```rust,no_run
8//! use mcp_kit::server::progress::ProgressTracker;
9//! use mcp_kit::server::NotificationSender;
10//! use mcp_kit::protocol::ProgressToken;
11//!
12//! async fn long_operation(
13//!     notifier: NotificationSender,
14//!     progress_token: Option<ProgressToken>,
15//! ) {
16//!     let tracker = ProgressTracker::new(notifier, progress_token);
17//!     
18//!     for i in 0..100 {
19//!         // Do some work...
20//!         tracker.update(i as f64, 100.0, Some(format!("Processing item {}", i))).await;
21//!     }
22//!     
23//!     tracker.complete("Done!").await;
24//! }
25//! ```
26
27use crate::protocol::ProgressToken;
28use crate::server::NotificationSender;
29
30/// A helper for sending progress updates for a long-running operation.
31///
32/// This wraps a `NotificationSender` and an optional progress token, making it
33/// easy to report progress without checking if the token is present each time.
34pub struct ProgressTracker {
35    notifier: NotificationSender,
36    token: Option<ProgressToken>,
37}
38
39impl ProgressTracker {
40    /// Create a new progress tracker.
41    ///
42    /// If `token` is `None`, all progress updates will be no-ops.
43    pub fn new(notifier: NotificationSender, token: Option<ProgressToken>) -> Self {
44        Self { notifier, token }
45    }
46
47    /// Create a tracker from a request's `_meta.progressToken`.
48    ///
49    /// Extracts the progress token from the `_meta` field of a request params object.
50    pub fn from_meta(notifier: NotificationSender, meta: Option<&serde_json::Value>) -> Self {
51        let token = meta.and_then(|m| m.get("progressToken")).and_then(|v| {
52            if let Some(s) = v.as_str() {
53                Some(ProgressToken::String(s.to_owned()))
54            } else {
55                v.as_i64().map(ProgressToken::Number)
56            }
57        });
58        Self::new(notifier, token)
59    }
60
61    /// Send a progress update.
62    ///
63    /// - `progress`: Current progress value
64    /// - `total`: Total value (progress/total gives percentage)
65    /// - `message`: Optional status message
66    ///
67    /// Returns silently if no progress token was provided.
68    pub async fn update(&self, progress: f64, total: f64, message: Option<String>) {
69        if let Some(ref token) = self.token {
70            let _ = self
71                .notifier
72                .progress(token.clone(), progress, Some(total), message)
73                .await;
74        }
75    }
76
77    /// Send a progress update with just a percentage (0.0 to 1.0).
78    pub async fn update_percent(&self, percent: f64, message: Option<String>) {
79        if let Some(ref token) = self.token {
80            let _ = self
81                .notifier
82                .progress(token.clone(), percent, None, message)
83                .await;
84        }
85    }
86
87    /// Send a progress update with a message.
88    pub async fn update_with_message(&self, progress: f64, total: f64, message: impl Into<String>) {
89        self.update(progress, total, Some(message.into())).await;
90    }
91
92    /// Mark the operation as complete with a final message.
93    pub async fn complete(&self, message: impl Into<String>) {
94        if let Some(ref token) = self.token {
95            let _ = self
96                .notifier
97                .progress(token.clone(), 1.0, Some(1.0), Some(message.into()))
98                .await;
99        }
100    }
101
102    /// Check if progress tracking is enabled (token was provided).
103    pub fn is_tracking(&self) -> bool {
104        self.token.is_some()
105    }
106}
107
108/// Extension trait to extract progress token from request parameters.
109pub trait ProgressTokenExt {
110    /// Get the progress token from `_meta.progressToken` if present.
111    fn progress_token(&self) -> Option<ProgressToken>;
112}
113
114impl ProgressTokenExt for serde_json::Value {
115    fn progress_token(&self) -> Option<ProgressToken> {
116        self.get("_meta")
117            .and_then(|m| m.get("progressToken"))
118            .and_then(|v| {
119                if let Some(s) = v.as_str() {
120                    Some(ProgressToken::String(s.to_owned()))
121                } else {
122                    v.as_i64().map(ProgressToken::Number)
123                }
124            })
125    }
126}