solid_mcp 0.2.2 → 0.5.0

data/ext/solid_mcp_native/core/src/pubsub.rs

@@ -0,0 +1,309 @@
+//! High-level pub/sub API
+//!
+//! This is the main interface for solid-mcp-core, providing:
+//! - Session-based subscriptions
+//! - Non-blocking message broadcasting
+//! - Graceful shutdown
+
+use crate::db::{Database, DbPool};
+use crate::subscriber::{MessageCallback, Subscriber};
+use crate::writer::MessageWriter;
+use crate::{Config, Error, Message, Result};
+use std::collections::HashMap;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+use tracing::{debug, info};
+
+/// The main pub/sub engine
+pub struct PubSub {
+    db: Arc<DbPool>,
+    config: Config,
+    writer: Arc<MessageWriter>,
+    subscribers: RwLock<HashMap<String, Subscriber>>,
+}
+
+impl PubSub {
+    /// Create a new pub/sub engine
+    pub async fn new(config: Config) -> Result<Self> {
+        let db = Arc::new(DbPool::new(&config).await?);
+        let writer = Arc::new(MessageWriter::new(db.clone(), &config).await?);
+
+        info!("PubSub engine initialized");
+
+        Ok(Self {
+            db,
+            config,
+            writer,
+            subscribers: RwLock::new(HashMap::new()),
+        })
+    }
+
+    /// Create a new pub/sub engine with an existing database pool
+    pub async fn with_db(db: Arc<DbPool>, config: Config) -> Result<Self> {
+        let writer = Arc::new(MessageWriter::new(db.clone(), &config).await?);
+
+        Ok(Self {
+            db,
+            config,
+            writer,
+            subscribers: RwLock::new(HashMap::new()),
+        })
+    }
+
+    /// Broadcast a message to a session (non-blocking)
+    ///
+    /// Returns `true` if the message was enqueued, `false` if the queue was full.
+    pub fn broadcast(
+        &self,
+        session_id: impl Into<String>,
+        event_type: impl Into<String>,
+        data: impl Into<String>,
+    ) -> Result<bool> {
+        let message = Message::new(session_id, event_type, data);
+        self.writer.enqueue(message)
+    }
+
+    /// Broadcast a message to a session (async, waits if queue is full)
+    pub async fn broadcast_async(
+        &self,
+        session_id: impl Into<String>,
+        event_type: impl Into<String>,
+        data: impl Into<String>,
+    ) -> Result<()> {
+        let message = Message::new(session_id, event_type, data);
+        self.writer.enqueue_async(message).await
+    }
+
+    /// Subscribe to messages for a session
+    ///
+    /// The callback will be invoked for each new message.
+    /// Returns an error if already subscribed to this session.
+    pub async fn subscribe(
+        &self,
+        session_id: impl Into<String>,
+        callback: MessageCallback,
+    ) -> Result<()> {
+        let session_id = session_id.into();
+
+        let mut subscribers = self.subscribers.write().await;
+
+        if subscribers.contains_key(&session_id) {
+            return Err(Error::Config(format!(
+                "Already subscribed to session {}",
+                session_id
+            )));
+        }
+
+        let subscriber =
+            Subscriber::new(&session_id, self.db.clone(), &self.config, callback).await?;
+        subscribers.insert(session_id, subscriber);
+
+        Ok(())
+    }
+
+    /// Unsubscribe from a session
+    pub async fn unsubscribe(&self, session_id: &str) -> Result<()> {
+        let mut subscribers = self.subscribers.write().await;
+
+        if let Some(subscriber) = subscribers.remove(session_id) {
+            subscriber.stop().await?;
+        }
+
+        Ok(())
+    }
+
+    /// Check if subscribed to a session
+    pub async fn is_subscribed(&self, session_id: &str) -> bool {
+        let subscribers = self.subscribers.read().await;
+        subscribers.contains_key(session_id)
+    }
+
+    /// Get the number of active subscriptions
+    pub async fn subscription_count(&self) -> usize {
+        let subscribers = self.subscribers.read().await;
+        subscribers.len()
+    }
+
+    /// Flush all pending messages to the database
+    pub async fn flush(&self) -> Result<()> {
+        self.writer.flush().await
+    }
+
+    /// Mark messages as delivered
+    pub async fn mark_delivered(&self, ids: &[i64]) -> Result<()> {
+        self.db.mark_delivered(ids).await
+    }
+
+    /// Cleanup old messages
+    pub async fn cleanup(&self) -> Result<(u64, u64)> {
+        let delivered = self
+            .db
+            .cleanup_delivered(self.config.delivered_retention)
+            .await?;
+        let undelivered = self
+            .db
+            .cleanup_undelivered(self.config.undelivered_retention)
+            .await?;
+        debug!(
+            "Cleanup complete: {} delivered, {} undelivered messages removed",
+            delivered, undelivered
+        );
+        Ok((delivered, undelivered))
+    }
+
+    /// Shutdown the pub/sub engine gracefully
+    pub async fn shutdown(self) -> Result<()> {
+        info!("PubSub engine shutting down...");
+
+        // Stop all subscribers
+        let mut subscribers = self.subscribers.write().await;
+        for (session_id, subscriber) in subscribers.drain() {
+            debug!("Stopping subscriber for session {}", session_id);
+            if let Err(e) = subscriber.stop().await {
+                tracing::error!("Error stopping subscriber for {}: {}", session_id, e);
+            }
+        }
+        drop(subscribers);
+
+        // Shutdown writer (flushes remaining messages)
+        // Need to unwrap the Arc - this only works if we're the last holder
+        match Arc::try_unwrap(self.writer) {
+            Ok(writer) => writer.shutdown().await?,
+            Err(_) => {
+                tracing::warn!("Could not unwrap writer Arc, forcing flush");
+                // Best effort - can't fully shutdown
+            }
+        }
+
+        info!("PubSub engine shutdown complete");
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::db::sqlite::SqlitePool;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::time::Duration;
+
+    async fn create_test_pubsub(config: Config) -> PubSub {
+        let sqlite = SqlitePool::new("sqlite::memory:").await.unwrap();
+        sqlite.setup_test_schema().await.unwrap();
+        let db = Arc::new(DbPool::Sqlite(sqlite));
+        PubSub::with_db(db, config).await.unwrap()
+    }
+
+    #[tokio::test]
+    async fn test_pubsub_basic() {
+        let config = Config::new("sqlite::memory:")
+            .batch_size(10)
+            .polling_interval(Duration::from_millis(10));
+
+        let pubsub = create_test_pubsub(config).await;
+
+        let received = Arc::new(AtomicUsize::new(0));
+        let received_clone = received.clone();
+
+        pubsub
+            .subscribe(
+                "session-1",
+                Box::new(move |_| {
+                    received_clone.fetch_add(1, Ordering::SeqCst);
+                }),
+            )
+            .await
+            .unwrap();
+
+        // Broadcast messages
+        for i in 0..5 {
+            pubsub
+                .broadcast("session-1", "message", format!(r#"{{"i":{}}}"#, i))
+                .unwrap();
+        }
+
+        // Flush and wait for delivery
+        pubsub.flush().await.unwrap();
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        assert_eq!(received.load(Ordering::SeqCst), 5);
+
+        pubsub.shutdown().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_pubsub_multiple_sessions() {
+        let config = Config::new("sqlite::memory:").polling_interval(Duration::from_millis(10));
+
+        let pubsub = create_test_pubsub(config).await;
+
+        let received1 = Arc::new(AtomicUsize::new(0));
+        let received2 = Arc::new(AtomicUsize::new(0));
+
+        let r1 = received1.clone();
+        let r2 = received2.clone();
+
+        pubsub
+            .subscribe(
+                "session-1",
+                Box::new(move |_| {
+                    r1.fetch_add(1, Ordering::SeqCst);
+                }),
+            )
+            .await
+            .unwrap();
+
+        pubsub
+            .subscribe(
+                "session-2",
+                Box::new(move |_| {
+                    r2.fetch_add(1, Ordering::SeqCst);
+                }),
+            )
+            .await
+            .unwrap();
+
+        // Broadcast to different sessions
+        pubsub.broadcast("session-1", "msg", "{}").unwrap();
+        pubsub.broadcast("session-1", "msg", "{}").unwrap();
+        pubsub.broadcast("session-2", "msg", "{}").unwrap();
+
+        pubsub.flush().await.unwrap();
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        assert_eq!(received1.load(Ordering::SeqCst), 2);
+        assert_eq!(received2.load(Ordering::SeqCst), 1);
+
+        pubsub.shutdown().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_pubsub_unsubscribe() {
+        let config = Config::new("sqlite::memory:").polling_interval(Duration::from_millis(10));
+
+        let pubsub = create_test_pubsub(config).await;
+
+        let received = Arc::new(AtomicUsize::new(0));
+        let r = received.clone();
+
+        pubsub
+            .subscribe(
+                "session-1",
+                Box::new(move |_| {
+                    r.fetch_add(1, Ordering::SeqCst);
+                }),
+            )
+            .await
+            .unwrap();
+
+        assert!(pubsub.is_subscribed("session-1").await);
+        assert_eq!(pubsub.subscription_count().await, 1);
+
+        pubsub.unsubscribe("session-1").await.unwrap();
+
+        assert!(!pubsub.is_subscribed("session-1").await);
+        assert_eq!(pubsub.subscription_count().await, 0);
+
+        pubsub.shutdown().await.unwrap();
+    }
+}

data/ext/solid_mcp_native/core/src/subscriber.rs

@@ -0,0 +1,298 @@
+//! Async subscriber for session-based message delivery
+//!
+//! Supports two modes:
+//! - PostgreSQL: Uses LISTEN/NOTIFY for real-time delivery (no polling)
+//! - SQLite: Falls back to efficient async polling
+
+#[cfg(feature = "postgres")]
+use crate::db::postgres::PostgresPool;
+use crate::db::{Database, DbPool};
+use crate::{Config, Message, Result};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
+use std::time::Duration;
+use tokio::task::JoinHandle;
+use tracing::{debug, error, info, warn};
+
+/// Callback type for message delivery
+pub type MessageCallback = Box<dyn Fn(Message) + Send + Sync + 'static>;
+
+/// A subscriber for a specific session
+pub struct Subscriber {
+    session_id: String,
+    handle: JoinHandle<()>,
+    shutdown: Arc<AtomicBool>,
+}
+
+impl Subscriber {
+    /// Create a new subscriber for a session
+    ///
+    /// The callback will be invoked for each new message.
+    pub async fn new(
+        session_id: impl Into<String>,
+        db: Arc<DbPool>,
+        config: &Config,
+        callback: MessageCallback,
+    ) -> Result<Self> {
+        let session_id = session_id.into();
+        let shutdown = Arc::new(AtomicBool::new(false));
+        let shutdown_clone = shutdown.clone();
+        let session_clone = session_id.clone();
+        let polling_interval = config.polling_interval;
+
+        // Get initial last_id
+        let last_id = Arc::new(AtomicI64::new(db.max_id().await?));
+
+        let handle = match &*db {
+            #[cfg(feature = "postgres")]
+            DbPool::Postgres(pg) => {
+                // Use LISTEN/NOTIFY for PostgreSQL
+                let pg_clone = pg.clone();
+                let db_clone = db.clone();
+                tokio::spawn(async move {
+                    postgres_subscriber_loop(
+                        session_clone,
+                        pg_clone,
+                        db_clone,
+                        last_id,
+                        shutdown_clone,
+                        callback,
+                    )
+                    .await
+                })
+            }
+            #[cfg(feature = "sqlite")]
+            DbPool::Sqlite(_) => {
+                // Use polling for SQLite
+                tokio::spawn(async move {
+                    polling_subscriber_loop(
+                        session_clone,
+                        db,
+                        last_id,
+                        polling_interval,
+                        shutdown_clone,
+                        callback,
+                    )
+                    .await
+                })
+            }
+        };
+
+        info!("Subscriber started for session {}", session_id);
+
+        Ok(Self {
+            session_id,
+            handle,
+            shutdown,
+        })
+    }
+
+    /// Stop the subscriber
+    pub async fn stop(self) -> Result<()> {
+        info!("Stopping subscriber for session {}", self.session_id);
+        self.shutdown.store(true, Ordering::SeqCst);
+
+        // Wait for the task to complete (with timeout)
+        tokio::select! {
+            _ = self.handle => {
+                debug!("Subscriber task completed");
+            }
+            _ = tokio::time::sleep(Duration::from_secs(5)) => {
+                warn!("Subscriber task did not complete in time, aborting");
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Get the session ID
+    pub fn session_id(&self) -> &str {
+        &self.session_id
+    }
+}
+
+/// Polling-based subscriber loop (for SQLite)
+async fn polling_subscriber_loop(
+    session_id: String,
+    db: Arc<DbPool>,
+    last_id: Arc<AtomicI64>,
+    polling_interval: Duration,
+    shutdown: Arc<AtomicBool>,
+    callback: MessageCallback,
+) {
+    debug!(
+        "Starting polling subscriber for session {} (interval: {:?})",
+        session_id, polling_interval
+    );
+
+    while !shutdown.load(Ordering::SeqCst) {
+        // Fetch new messages
+        let current_last_id = last_id.load(Ordering::SeqCst);
+        match db.fetch_after(&session_id, current_last_id, 100).await {
+            Ok(messages) => {
+                for msg in messages {
+                    let msg_id = msg.id;
+
+                    // Deliver to callback
+                    callback(msg);
+
+                    // Update last_id
+                    last_id.store(msg_id, Ordering::SeqCst);
+                }
+            }
+            Err(e) => {
+                error!("Error fetching messages for session {}: {}", session_id, e);
+            }
+        }
+
+        // Sleep until next poll (interruptible)
+        tokio::select! {
+            _ = tokio::time::sleep(polling_interval) => {}
+            _ = async {
+                while !shutdown.load(Ordering::SeqCst) {
+                    tokio::time::sleep(Duration::from_millis(10)).await;
+                }
+            } => {
+                break;
+            }
+        }
+    }
+
+    debug!("Polling subscriber for session {} stopped", session_id);
+}
+
+/// LISTEN/NOTIFY-based subscriber loop (for PostgreSQL)
+#[cfg(feature = "postgres")]
+async fn postgres_subscriber_loop(
+    session_id: String,
+    pg: PostgresPool,
+    db: Arc<DbPool>,
+    last_id: Arc<AtomicI64>,
+    shutdown: Arc<AtomicBool>,
+    callback: MessageCallback,
+) {
+    debug!(
+        "Starting LISTEN/NOTIFY subscriber for session {}",
+        session_id
+    );
+
+    // First, catch up on any missed messages
+    let current_last_id = last_id.load(Ordering::SeqCst);
+    match db.fetch_after(&session_id, current_last_id, 1000).await {
+        Ok(messages) => {
+            for msg in messages {
+                let msg_id = msg.id;
+                callback(msg);
+                last_id.store(msg_id, Ordering::SeqCst);
+            }
+        }
+        Err(e) => {
+            error!(
+                "Error catching up messages for session {}: {}",
+                session_id, e
+            );
+        }
+    }
+
+    // Set up LISTEN
+    let mut listener = match pg.listen(&session_id).await {
+        Ok(l) => l,
+        Err(e) => {
+            error!(
+                "Failed to create listener for session {}: {}",
+                session_id, e
+            );
+            return;
+        }
+    };
+
+    // Listen for notifications
+    while !shutdown.load(Ordering::SeqCst) {
+        tokio::select! {
+            notification = listener.recv() => {
+                match notification {
+                    Ok(notif) => {
+                        // Notification payload is the message ID
+                        if let Ok(msg_id) = notif.payload().parse::<i64>() {
+                            // Fetch the specific message
+                            let current = last_id.load(Ordering::SeqCst);
+                            if msg_id > current {
+                                match db.fetch_after(&session_id, current, 100).await {
+                                    Ok(messages) => {
+                                        for msg in messages {
+                                            let id = msg.id;
+                                            callback(msg);
+                                            last_id.store(id, Ordering::SeqCst);
+                                        }
+                                    }
+                                    Err(e) => {
+                                        error!("Error fetching message {}: {}", msg_id, e);
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        error!("Listener error for session {}: {}", session_id, e);
+                        // Reconnect logic could go here
+                        break;
+                    }
+                }
+            }
+            _ = tokio::time::sleep(Duration::from_secs(1)) => {
+                // Periodic check for shutdown
+                if shutdown.load(Ordering::SeqCst) {
+                    break;
+                }
+            }
+        }
+    }
+
+    debug!(
+        "LISTEN/NOTIFY subscriber for session {} stopped",
+        session_id
+    );
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::db::sqlite::SqlitePool;
+    use std::sync::atomic::AtomicUsize;
+
+    async fn create_test_db() -> Arc<DbPool> {
+        let sqlite = SqlitePool::new("sqlite::memory:").await.unwrap();
+        sqlite.setup_test_schema().await.unwrap();
+        Arc::new(DbPool::Sqlite(sqlite))
+    }
+
+    #[tokio::test]
+    async fn test_polling_subscriber() {
+        let db = create_test_db().await;
+        let config = Config::new("sqlite::memory:").polling_interval(Duration::from_millis(10));
+
+        let received = Arc::new(AtomicUsize::new(0));
+        let received_clone = received.clone();
+
+        let callback: MessageCallback = Box::new(move |_msg| {
+            received_clone.fetch_add(1, Ordering::SeqCst);
+        });
+
+        let subscriber = Subscriber::new("session-1", db.clone(), &config, callback)
+            .await
+            .unwrap();
+
+        // Insert some messages
+        let messages: Vec<Message> = (0..5)
+            .map(|i| Message::new("session-1", "message", format!(r#"{{"i":{}}}"#, i)))
+            .collect();
+        db.insert_batch(&messages).await.unwrap();
+
+        // Wait for subscriber to pick them up
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        assert_eq!(received.load(Ordering::SeqCst), 5);
+
+        subscriber.stop().await.unwrap();
+    }
+}