mcp_host/managers/

progress.rs

//! MCP progress tracking for long-running operations.
//!
//! Provides progress notifications with unique tokens, tracking operation state,
//! and automatic cleanup of stale tokens. Tokens are formatted as:
//! `"{op_type}-{unix_timestamp}-{counter}"`

use dashmap::DashMap;
use serde_json::json;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::mpsc;

use crate::transport::traits::JsonRpcNotification;

/// Tracks state of an active operation
#[derive(Debug, Clone)]
struct OperationState {
    /// Operation type: "stalint", "dictator", etc.
    #[allow(dead_code)]
    op_type: String,
    /// Stable operation ID (same as token)
    #[allow(dead_code)]
    op_id: String,
    /// Total items to process
    total: u32,
    /// Currently processed items
    current: u32,
    /// When operation started
    start_time: Instant,
}

/// Progress token for tracking an operation
pub type ProgressToken = String;

/// Progress tracker for long-running MCP operations
///
/// # Example
///
/// ```rust,ignore
/// let tracker = ProgressTracker::new(notification_tx);
/// let token = tracker.start("my-operation", 100);
/// tracker.progress(&token, 50);  // 50% complete
/// tracker.finish(&token);        // 100% complete
/// ```
#[derive(Clone)]
pub struct ProgressTracker {
    /// Global counter for unique progress tokens
    token_counter: Arc<AtomicU64>,
    /// Active operations: token -> state
    operations: Arc<DashMap<String, OperationState>>,
    /// Notification channel sender
    notif_tx: mpsc::UnboundedSender<JsonRpcNotification>,
}

impl ProgressTracker {
    /// Create a new progress tracker
    ///
    /// # Arguments
    ///
    /// * `notif_tx` - Notification channel for sending progress updates
    pub fn new(notif_tx: mpsc::UnboundedSender<JsonRpcNotification>) -> Self {
        Self {
            token_counter: Arc::new(AtomicU64::new(0)),
            operations: Arc::new(DashMap::new()),
            notif_tx,
        }
    }

    /// Generate a unique progress token: "{op_type}-{timestamp}-{counter}"
    fn generate_token(&self, op_type: &str) -> ProgressToken {
        let counter = self.token_counter.fetch_add(1, Ordering::SeqCst);
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        format!("{op_type}-{timestamp}-{counter}")
    }

    /// Start a new operation and return its unique progress token
    ///
    /// Sends initial progress notification (0%)
    ///
    /// # Arguments
    ///
    /// * `op_type` - Operation type identifier (e.g., "lint", "fix")
    /// * `total` - Total number of items to process
    ///
    /// # Returns
    ///
    /// Unique progress token to use for updates
    pub fn start(&self, op_type: &str, total: u32) -> ProgressToken {
        let token = self.generate_token(op_type);
        let op_id = token.clone();

        let op_state = OperationState {
            op_type: op_type.to_string(),
            op_id,
            total,
            current: 0,
            start_time: Instant::now(),
        };

        self.operations.insert(token.clone(), op_state);

        // Send initial progress notification (0%)
        self.send_notification(&token, 0, total);

        token
    }

    /// Update progress for an operation (0..=total)
    ///
    /// # Arguments
    ///
    /// * `token` - Progress token from `start()`
    /// * `current` - Current progress (will be clamped to [0, total])
    ///
    /// # Returns
    ///
    /// `true` if operation is still active, `false` if unknown/expired
    pub fn progress(&self, token: &ProgressToken, current: u32) -> bool {
        if let Some(mut op) = self.operations.get_mut(token) {
            // Clamp progress to [0, total]
            op.current = current.min(op.total);
            let progress = op.current;
            let total = op.total;
            drop(op); // Release lock before sending notification

            self.send_notification(token, progress, total);
            true
        } else {
            false
        }
    }

    /// Mark operation complete (final notification at 100%)
    ///
    /// Removes the operation from tracking and sends final progress notification.
    ///
    /// # Arguments
    ///
    /// * `token` - Progress token from `start()`
    pub fn finish(&self, token: &ProgressToken) {
        if let Some((_key, op)) = self.operations.remove(token) {
            self.send_notification(token, op.total, op.total);
        }
    }

    /// Send progress notification to client
    fn send_notification(&self, token: &ProgressToken, progress: u32, total: u32) {
        let notification = JsonRpcNotification::new(
            "notifications/progress",
            Some(json!({
                "progressToken": token,
                "progress": progress,
                "total": total
            })),
        );

        let _ = self.notif_tx.send(notification);
    }

    /// Clean up stale operations (older than 10 minutes)
    ///
    /// Call periodically from background tasks. Only performs cleanup
    /// if at least 60 seconds have passed since last cleanup.
    pub fn cleanup_stale(&self) {
        let now = Instant::now();
        let timeout = Duration::from_secs(600); // 10 minutes

        self.operations
            .retain(|_token, op| now.duration_since(op.start_time) < timeout);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_tracker() -> ProgressTracker {
        let (tx, _rx) = mpsc::unbounded_channel();
        ProgressTracker::new(tx)
    }

    #[test]
    fn test_token_format() {
        let tracker = create_tracker();
        let token = tracker.generate_token("stalint");

        // Format: "{op_type}-{timestamp}-{counter}"
        let parts: Vec<&str> = token.split('-').collect();
        assert_eq!(parts.len(), 3);
        assert_eq!(parts[0], "stalint");
        // parts[1] is unix timestamp
        assert!(parts[1].parse::<u64>().is_ok());
        // parts[2] is counter
        assert!(parts[2].parse::<u64>().is_ok());
    }

    #[test]
    fn test_token_uniqueness() {
        let tracker = create_tracker();
        let token1 = tracker.generate_token("stalint");
        let token2 = tracker.generate_token("stalint");

        assert_ne!(token1, token2);
    }

    #[test]
    fn test_start_registers_operation() {
        let tracker = create_tracker();
        let token = tracker.start("stalint", 100);

        assert!(tracker.operations.contains_key(&token));
        let op = tracker.operations.get(&token).unwrap();
        assert_eq!(op.current, 0);
        assert_eq!(op.total, 100);
    }

    #[test]
    fn test_progress_updates() {
        let tracker = create_tracker();
        let token = tracker.start("dictator", 50);

        // Update to 50%
        assert!(tracker.progress(&token, 25));
        let op = tracker.operations.get(&token).unwrap();
        assert_eq!(op.current, 25);
        assert_eq!(op.total, 50);
        drop(op);

        // Update to 100%
        assert!(tracker.progress(&token, 50));
        let op = tracker.operations.get(&token).unwrap();
        assert_eq!(op.current, 50);
        assert_eq!(op.total, 50);
    }

    #[test]
    fn test_progress_clamps_to_total() {
        let tracker = create_tracker();
        let token = tracker.start("supremecourt", 10);

        // Try to set beyond total
        assert!(tracker.progress(&token, 99));
        let op = tracker.operations.get(&token).unwrap();
        assert_eq!(op.current, 10); // Clamped to total
        assert_eq!(op.total, 10);
    }

    #[test]
    fn test_finish_removes_operation() {
        let tracker = create_tracker();
        let token = tracker.start("stalint", 100);

        assert!(tracker.operations.contains_key(&token));
        tracker.finish(&token);
        assert!(!tracker.operations.contains_key(&token));
    }

    #[test]
    fn test_progress_unknown_token() {
        let tracker = create_tracker();

        // Unknown token should return false
        assert!(!tracker.progress(&"unknown-token".to_string(), 50));
    }

    #[test]
    fn test_multiple_concurrent_operations() {
        let tracker = create_tracker();
        let token1 = tracker.start("stalint", 100);
        let token2 = tracker.start("dictator", 50);

        assert!(tracker.operations.contains_key(&token1));
        assert!(tracker.operations.contains_key(&token2));

        // Update both independently
        assert!(tracker.progress(&token1, 50));
        assert!(tracker.progress(&token2, 25));

        let op1 = tracker.operations.get(&token1).unwrap();
        let op2 = tracker.operations.get(&token2).unwrap();

        assert_eq!((op1.current, op1.total), (50, 100));
        assert_eq!((op2.current, op2.total), (25, 50));
        drop(op1);
        drop(op2);

        // Finish one
        tracker.finish(&token1);
        assert!(!tracker.operations.contains_key(&token1));
        assert!(tracker.operations.contains_key(&token2));
    }

    #[test]
    fn test_cleanup_stale() {
        let tracker = create_tracker();
        let token = tracker.start("test", 100);

        assert!(tracker.operations.contains_key(&token));

        // Manually age the operation
        if let Some(mut op) = tracker.operations.get_mut(&token) {
            op.start_time = Instant::now() - Duration::from_secs(700); // 11+ minutes ago
        }

        tracker.cleanup_stale();

        // Should be removed
        assert!(!tracker.operations.contains_key(&token));
    }
}