remote-mcp-kernel 0.1.0-alpha.6

//! SSE (Server-Sent Events) handler for MCP communication
//!
//! This module provides an independent SSE handler for MCP protocol
//! communication, following microkernel design principles where
//! each handler has a single, well-defined responsibility.

use axum::{Router, middleware};
use rmcp::transport::sse_server::{SseServer, SseServerConfig};
use tokio_util::sync::CancellationToken;

use crate::{error::AppResult, handlers::McpServerHandler};
use oauth_provider_rs::http_integration::middleware::simple_auth_middleware;

/// SSE handler for MCP protocol communication
///
/// This handler is responsible solely for managing SSE connections
/// and MCP protocol communication. It's designed to be independent
/// and composable following microkernel architecture principles.
#[derive(Clone)]
pub struct SseHandler<M: McpServerHandler> {
    /// MCP server (dependency injection)
    mcp_server: M,
}

/// SSE handler configuration
#[derive(Debug, Clone)]
pub struct SseHandlerConfig {
    pub sse_path: String,
    pub message_path: String,
    pub keep_alive_seconds: u64,
    pub require_auth: bool,
}

impl Default for SseHandlerConfig {
    fn default() -> Self {
        Self {
            sse_path: "/mcp/sse".to_string(),
            message_path: "/mcp/message".to_string(),
            keep_alive_seconds: 15,
            require_auth: true,
        }
    }
}

impl SseHandlerConfig {
    /// Get default configuration for SSE handler
    pub fn default_config() -> SseHandlerConfig {
        SseHandlerConfig::default()
    }

    /// Create SSE handler configuration with custom paths
    pub fn config_with_paths(
        sse_path: impl Into<String>,
        message_path: impl Into<String>,
    ) -> SseHandlerConfig {
        SseHandlerConfig {
            sse_path: sse_path.into(),
            message_path: message_path.into(),
            ..Default::default()
        }
    }

    /// Create SSE handler configuration without authentication
    pub fn config_without_auth() -> SseHandlerConfig {
        SseHandlerConfig {
            require_auth: false,
            ..Default::default()
        }
    }
}

impl<M: McpServerHandler> SseHandler<M> {
    /// Create a new SSE handler with MCP server
    pub fn new(mcp_server: M) -> Self {
        Self { mcp_server }
    }

    /// Create router for SSE endpoints
    ///
    /// Returns a router with SSE endpoints that can be composed
    /// with other routers following microkernel principles.
    pub fn router(&self, config: SseHandlerConfig) -> AppResult<Router> {
        let sse_config = SseServerConfig {
            bind: "0.0.0.0:0".parse().unwrap(), // Will be overridden by actual server
            sse_path: config.sse_path,
            post_path: config.message_path,
            ct: CancellationToken::new(),
            sse_keep_alive: Some(std::time::Duration::from_secs(config.keep_alive_seconds)),
        };

        let (sse_server, sse_router) = SseServer::new(sse_config);

        // Register MCP service with the SSE server
        let mcp_server = self.mcp_server.clone();
        let _service_token = sse_server.with_service(move || mcp_server.clone());

        // Apply authentication middleware if required
        let router = if config.require_auth {
            sse_router.layer(middleware::from_fn(simple_auth_middleware))
        } else {
            sse_router
        };

        Ok(router)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::handlers::McpServer;
    use oauth_provider_rs::{
        GitHubOAuthProvider, OAuthProvider, provider_trait::OAuthProviderConfig,
    };

    #[test]
    fn test_sse_handler_creation() {
        let github_config = OAuthProviderConfig::with_oauth_config(
            "test_client_id".to_string(),
            "test_client_secret".to_string(),
            "http://localhost:8080/oauth/callback".to_string(),
            "read:user".to_string(),
            "github".to_string(),
        );

        let oauth_provider = OAuthProvider::new(
            GitHubOAuthProvider::new_github(github_config),
            oauth_provider_rs::http_integration::config::OAuthProviderConfig::default(),
        );
        let mcp_server = McpServer::new();
        let _sse_handler = SseHandler::new(mcp_server);

        // Handler should be created successfully
    }

    #[test]
    fn test_sse_config_defaults() {
        let config = SseHandlerConfig::default();

        assert_eq!(config.sse_path, "/mcp/sse");
        assert_eq!(config.message_path, "/mcp/message");
        assert_eq!(config.keep_alive_seconds, 15);
        assert!(config.require_auth);
    }

    #[test]
    fn test_sse_config_custom_paths() {
        let config = SseHandlerConfig::config_with_paths("/custom/sse", "/custom/message");

        assert_eq!(config.sse_path, "/custom/sse");
        assert_eq!(config.message_path, "/custom/message");
        assert!(config.require_auth);
    }

    #[test]
    fn test_sse_config_without_auth() {
        let config = SseHandlerConfig::config_without_auth();

        assert!(!config.require_auth);
        assert_eq!(config.sse_path, "/mcp/sse");
    }

    #[tokio::test]
    async fn test_sse_router_creation() {
        let github_config = OAuthProviderConfig::with_oauth_config(
            "test_client_id".to_string(),
            "test_client_secret".to_string(),
            "http://localhost:8080/oauth/callback".to_string(),
            "read:user".to_string(),
            "github".to_string(),
        );

        let oauth_provider = OAuthProvider::new(
            GitHubOAuthProvider::new_github(github_config),
            oauth_provider_rs::http_integration::config::OAuthProviderConfig::default(),
        );
        let mcp_server = McpServer::new();
        let sse_handler = SseHandler::new(mcp_server);
        let config = SseHandlerConfig::default();

        let router_result = sse_handler.router(config);
        assert!(router_result.is_ok());
    }
}