mcp_host/managers/

subscription.rs

//! Resource subscription manager
//!
//! Tracks which sessions are subscribed to which resource URIs and provides
//! notification dispatch when resources change.

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};

use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

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

/// Resource subscription request parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubscribeRequest {
    /// URI of the resource to subscribe to
    pub uri: String,
}

/// Resource unsubscribe request parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UnsubscribeRequest {
    /// URI of the resource to unsubscribe from
    pub uri: String,
}

/// Resource updated notification parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceUpdatedNotification {
    /// URI of the updated resource
    pub uri: String,
}

/// Subscription manager for tracking resource subscriptions
///
/// Thread-safe structure that maps resource URIs to subscribed sessions
/// and sessions to their subscribed URIs.
#[derive(Clone)]
pub struct SubscriptionManager {
    /// Maps resource URI -> set of session IDs
    uri_to_sessions: Arc<RwLock<HashMap<String, HashSet<String>>>>,

    /// Maps session ID -> set of resource URIs (for efficient cleanup)
    session_to_uris: Arc<RwLock<HashMap<String, HashSet<String>>>>,

    /// Notification channel for dispatching updates
    notification_tx: Option<mpsc::UnboundedSender<JsonRpcNotification>>,

    /// Per-session notification channels (for targeted notifications)
    session_notifiers: Arc<RwLock<HashMap<String, mpsc::UnboundedSender<JsonRpcNotification>>>>,
}

impl SubscriptionManager {
    /// Create a new subscription manager
    pub fn new() -> Self {
        Self {
            uri_to_sessions: Arc::new(RwLock::new(HashMap::new())),
            session_to_uris: Arc::new(RwLock::new(HashMap::new())),
            notification_tx: None,
            session_notifiers: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Create subscription manager with a global notification channel
    pub fn with_notifications(notification_tx: mpsc::UnboundedSender<JsonRpcNotification>) -> Self {
        Self {
            uri_to_sessions: Arc::new(RwLock::new(HashMap::new())),
            session_to_uris: Arc::new(RwLock::new(HashMap::new())),
            notification_tx: Some(notification_tx),
            session_notifiers: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a per-session notification channel
    ///
    /// This allows targeted notifications to specific sessions rather than
    /// broadcasting to all.
    pub fn register_session_notifier(
        &self,
        session_id: &str,
        tx: mpsc::UnboundedSender<JsonRpcNotification>,
    ) {
        if let Ok(mut notifiers) = self.session_notifiers.write() {
            notifiers.insert(session_id.to_string(), tx);
        }
    }

    /// Unregister a session's notification channel
    pub fn unregister_session_notifier(&self, session_id: &str) {
        if let Ok(mut notifiers) = self.session_notifiers.write() {
            notifiers.remove(session_id);
        }
    }

    /// Subscribe a session to a resource URI
    ///
    /// Returns true if this is a new subscription, false if already subscribed.
    pub fn subscribe(&self, session_id: &str, uri: &str) -> bool {
        let is_new_uri_subscription;

        // Add to uri_to_sessions
        {
            let mut uri_map = self.uri_to_sessions.write().unwrap();
            let sessions = uri_map.entry(uri.to_string()).or_default();
            is_new_uri_subscription = sessions.insert(session_id.to_string());
        }

        // Add to session_to_uris
        {
            let mut session_map = self.session_to_uris.write().unwrap();
            let uris = session_map.entry(session_id.to_string()).or_default();
            uris.insert(uri.to_string());
        }

        tracing::debug!(
            session_id = %session_id,
            uri = %uri,
            new = is_new_uri_subscription,
            "Resource subscription"
        );

        is_new_uri_subscription
    }

    /// Unsubscribe a session from a resource URI
    ///
    /// Returns true if subscription existed and was removed.
    pub fn unsubscribe(&self, session_id: &str, uri: &str) -> bool {
        let was_subscribed;

        // Remove from uri_to_sessions
        {
            let mut uri_map = self.uri_to_sessions.write().unwrap();
            if let Some(sessions) = uri_map.get_mut(uri) {
                was_subscribed = sessions.remove(session_id);
                // Clean up empty sets
                if sessions.is_empty() {
                    uri_map.remove(uri);
                }
            } else {
                was_subscribed = false;
            }
        }

        // Remove from session_to_uris
        {
            let mut session_map = self.session_to_uris.write().unwrap();
            if let Some(uris) = session_map.get_mut(session_id) {
                uris.remove(uri);
                // Clean up empty sets
                if uris.is_empty() {
                    session_map.remove(session_id);
                }
            }
        }

        tracing::debug!(
            session_id = %session_id,
            uri = %uri,
            was_subscribed = was_subscribed,
            "Resource unsubscription"
        );

        was_subscribed
    }

    /// Remove all subscriptions for a session
    ///
    /// Call this when a session closes to clean up all its subscriptions.
    pub fn remove_session(&self, session_id: &str) {
        // Get all URIs this session was subscribed to
        let uris: Vec<String> = {
            let mut session_map = self.session_to_uris.write().unwrap();
            session_map
                .remove(session_id)
                .map(|uris| uris.into_iter().collect())
                .unwrap_or_default()
        };

        // Remove session from each URI's subscriber set
        {
            let mut uri_map = self.uri_to_sessions.write().unwrap();
            for uri in &uris {
                if let Some(sessions) = uri_map.get_mut(uri) {
                    sessions.remove(session_id);
                    if sessions.is_empty() {
                        uri_map.remove(uri);
                    }
                }
            }
        }

        // Remove session notifier
        self.unregister_session_notifier(session_id);

        if !uris.is_empty() {
            tracing::debug!(
                session_id = %session_id,
                subscription_count = uris.len(),
                "Removed all session subscriptions"
            );
        }
    }

    /// Check if a session is subscribed to a URI
    pub fn is_subscribed(&self, session_id: &str, uri: &str) -> bool {
        self.uri_to_sessions
            .read()
            .ok()
            .and_then(|map| map.get(uri).map(|sessions| sessions.contains(session_id)))
            .unwrap_or(false)
    }

    /// Get all sessions subscribed to a URI
    pub fn get_subscribers(&self, uri: &str) -> Vec<String> {
        self.uri_to_sessions
            .read()
            .ok()
            .and_then(|map| map.get(uri).map(|sessions| sessions.iter().cloned().collect()))
            .unwrap_or_default()
    }

    /// Get all URIs a session is subscribed to
    pub fn get_session_subscriptions(&self, session_id: &str) -> Vec<String> {
        self.session_to_uris
            .read()
            .ok()
            .and_then(|map| map.get(session_id).map(|uris| uris.iter().cloned().collect()))
            .unwrap_or_default()
    }

    /// Get total subscription count
    pub fn subscription_count(&self) -> usize {
        self.uri_to_sessions
            .read()
            .ok()
            .map(|map| map.values().map(|s| s.len()).sum())
            .unwrap_or(0)
    }

    /// Get count of unique URIs with subscribers
    pub fn subscribed_uri_count(&self) -> usize {
        self.uri_to_sessions
            .read()
            .ok()
            .map(|map| map.len())
            .unwrap_or(0)
    }

    /// Notify all subscribers that a resource has been updated
    ///
    /// Sends `notifications/resources/updated` to all sessions subscribed to the URI.
    pub fn notify_resource_updated(&self, uri: &str) {
        let subscribers = self.get_subscribers(uri);

        if subscribers.is_empty() {
            return;
        }

        let notification = JsonRpcNotification::new(
            "notifications/resources/updated",
            Some(serde_json::json!({ "uri": uri })),
        );

        // Try per-session notifications first
        let notifiers = self.session_notifiers.read().ok();
        let mut notified_count = 0;

        if let Some(notifiers) = notifiers {
            for session_id in &subscribers {
                if let Some(tx) = notifiers.get(session_id)
                    && tx.send(notification.clone()).is_ok()
                {
                    notified_count += 1;
                }
            }
        }

        // Fall back to global notification channel if no per-session notifiers
        if notified_count == 0
            && let Some(ref tx) = self.notification_tx
        {
            // Note: This broadcasts to all connections, not just subscribers
            // In a multi-session scenario, you'd need session-aware transport
            let _ = tx.send(notification);
        }

        tracing::debug!(
            uri = %uri,
            subscriber_count = subscribers.len(),
            notified_count = notified_count,
            "Resource update notification sent"
        );
    }

    /// Notify subscribers for multiple URIs (batch operation)
    pub fn notify_resources_updated(&self, uris: &[&str]) {
        for uri in uris {
            self.notify_resource_updated(uri);
        }
    }
}

impl Default for SubscriptionManager {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_subscription_manager_creation() {
        let manager = SubscriptionManager::new();
        assert_eq!(manager.subscription_count(), 0);
        assert_eq!(manager.subscribed_uri_count(), 0);
    }

    #[test]
    fn test_subscribe() {
        let manager = SubscriptionManager::new();

        // First subscription returns true
        assert!(manager.subscribe("session-1", "file:///test.txt"));

        // Duplicate subscription returns false
        assert!(!manager.subscribe("session-1", "file:///test.txt"));

        // Different session, same URI returns true
        assert!(manager.subscribe("session-2", "file:///test.txt"));

        assert_eq!(manager.subscription_count(), 2);
        assert_eq!(manager.subscribed_uri_count(), 1);
    }

    #[test]
    fn test_unsubscribe() {
        let manager = SubscriptionManager::new();

        manager.subscribe("session-1", "file:///test.txt");
        manager.subscribe("session-2", "file:///test.txt");

        // Unsubscribe existing returns true
        assert!(manager.unsubscribe("session-1", "file:///test.txt"));

        // Unsubscribe again returns false
        assert!(!manager.unsubscribe("session-1", "file:///test.txt"));

        // Other session still subscribed
        assert!(manager.is_subscribed("session-2", "file:///test.txt"));

        assert_eq!(manager.subscription_count(), 1);
    }

    #[test]
    fn test_remove_session() {
        let manager = SubscriptionManager::new();

        manager.subscribe("session-1", "file:///a.txt");
        manager.subscribe("session-1", "file:///b.txt");
        manager.subscribe("session-2", "file:///a.txt");

        manager.remove_session("session-1");

        assert!(!manager.is_subscribed("session-1", "file:///a.txt"));
        assert!(!manager.is_subscribed("session-1", "file:///b.txt"));
        assert!(manager.is_subscribed("session-2", "file:///a.txt"));

        // file:///b.txt should be cleaned up since no subscribers
        assert_eq!(manager.subscribed_uri_count(), 1);
    }

    #[test]
    fn test_get_subscribers() {
        let manager = SubscriptionManager::new();

        manager.subscribe("session-1", "file:///test.txt");
        manager.subscribe("session-2", "file:///test.txt");
        manager.subscribe("session-3", "file:///other.txt");

        let subscribers = manager.get_subscribers("file:///test.txt");
        assert_eq!(subscribers.len(), 2);
        assert!(subscribers.contains(&"session-1".to_string()));
        assert!(subscribers.contains(&"session-2".to_string()));
    }

    #[test]
    fn test_get_session_subscriptions() {
        let manager = SubscriptionManager::new();

        manager.subscribe("session-1", "file:///a.txt");
        manager.subscribe("session-1", "file:///b.txt");
        manager.subscribe("session-2", "file:///c.txt");

        let subs = manager.get_session_subscriptions("session-1");
        assert_eq!(subs.len(), 2);
        assert!(subs.contains(&"file:///a.txt".to_string()));
        assert!(subs.contains(&"file:///b.txt".to_string()));
    }

    #[test]
    fn test_is_subscribed() {
        let manager = SubscriptionManager::new();

        assert!(!manager.is_subscribed("session-1", "file:///test.txt"));

        manager.subscribe("session-1", "file:///test.txt");

        assert!(manager.is_subscribed("session-1", "file:///test.txt"));
        assert!(!manager.is_subscribed("session-2", "file:///test.txt"));
    }

    #[tokio::test]
    async fn test_notify_resource_updated() {
        let (tx, mut rx) = mpsc::unbounded_channel();
        let manager = SubscriptionManager::with_notifications(tx);

        manager.subscribe("session-1", "file:///test.txt");

        manager.notify_resource_updated("file:///test.txt");

        // Should receive notification
        let notification = rx.recv().await.expect("Should receive notification");
        assert_eq!(notification.method, "notifications/resources/updated");

        let params = notification.params.expect("Should have params");
        assert_eq!(params["uri"], "file:///test.txt");
    }

    #[tokio::test]
    async fn test_per_session_notifier() {
        let manager = SubscriptionManager::new();
        let (tx1, mut rx1) = mpsc::unbounded_channel();
        let (tx2, mut rx2) = mpsc::unbounded_channel();

        manager.register_session_notifier("session-1", tx1);
        manager.register_session_notifier("session-2", tx2);

        manager.subscribe("session-1", "file:///test.txt");
        // session-2 not subscribed

        manager.notify_resource_updated("file:///test.txt");

        // session-1 should receive notification
        let notification = rx1.recv().await.expect("Session-1 should receive notification");
        assert_eq!(notification.method, "notifications/resources/updated");

        // session-2 should NOT receive (not subscribed)
        // Use try_recv to check without blocking
        assert!(rx2.try_recv().is_err());
    }
}