remote_mcp_kernel/microkernel/

core.rs

//! Core microkernel server implementation
//!
//! This module contains the main MicrokernelServer implementation
//! demonstrating the microkernel design pattern.

use axum::Router;
use std::net::SocketAddr;

use crate::{
    error::AppResult,
    handlers::{McpServerHandler, SseHandler, SseHandlerConfig, StreamableHttpHandler},
};
use oauth_provider_rs::{DefaultClientManager, OAuthProvider, OAuthProviderTrait, OAuthStorage};

/// Configuration for custom router attachment
#[derive(Debug, Clone)]
pub struct CustomRouterConfig {
    /// Optional name for the custom router (for logging/debugging)
    pub name: Option<String>,
    /// Optional path prefix for the router
    pub path_prefix: Option<String>,
}

impl Default for CustomRouterConfig {
    fn default() -> Self {
        Self {
            name: None,
            path_prefix: None,
        }
    }
}

/// Container for custom router with its configuration
pub struct CustomRouter {
    router: Router,
    config: CustomRouterConfig,
}

/// Microkernel server builder that composes independent handlers
///
/// This builder demonstrates the microkernel principle where services
/// are composed from independent, single-responsibility components.
/// Now supports any OAuth provider, storage backend, MCP server, and custom routers through trait abstraction.
pub struct MicrokernelServer<
    P: OAuthProviderTrait<S, DefaultClientManager<S>> + 'static,
    S: OAuthStorage + Clone + 'static,
    M: McpServerHandler = crate::handlers::McpServer,
> {
    oauth_provider: Option<OAuthProvider<P, S>>,
    streamable_handler: Option<StreamableHttpHandler<M>>,
    sse_handler: Option<SseHandler<M>>,
    sse_config: SseHandlerConfig,
    custom_routers: Vec<CustomRouter>,
}

impl<
    P: OAuthProviderTrait<S, DefaultClientManager<S>> + 'static,
    S: OAuthStorage + Clone + 'static,
    M: McpServerHandler,
> MicrokernelServer<P, S, M>
{
    /// Create a new microkernel server builder
    pub fn new() -> Self {
        Self {
            oauth_provider: None,
            streamable_handler: None,
            sse_handler: None,
            sse_config: SseHandlerConfig::default(),
            custom_routers: Vec::new(),
        }
    }

    /// Add OAuth provider handler
    pub fn with_oauth_provider(mut self, oauth_provider: OAuthProvider<P, S>) -> Self {
        self.oauth_provider = Some(oauth_provider);
        self
    }

    /// Add streamable HTTP handler
    pub fn with_streamable_handler(mut self, streamable_handler: StreamableHttpHandler<M>) -> Self {
        self.streamable_handler = Some(streamable_handler);
        self
    }

    /// Add SSE handler with configuration
    pub fn with_sse_handler(
        mut self,
        sse_handler: SseHandler<M>,
        config: SseHandlerConfig,
    ) -> Self {
        self.sse_handler = Some(sse_handler);
        self.sse_config = config;
        self
    }

    /// Create SSE handler with MCP server and configuration
    pub fn with_mcp_sse_handler(mut self, mcp_server: M, config: SseHandlerConfig) -> Self {
        let sse_handler = SseHandler::new(mcp_server);
        self.sse_handler = Some(sse_handler);
        self.sse_config = config;
        self
    }

    /// Create streamable HTTP handler with MCP server
    pub fn with_mcp_streamable_handler(mut self, mcp_server: M) -> Self {
        let streamable_handler = StreamableHttpHandler::new(mcp_server);
        self.streamable_handler = Some(streamable_handler);
        self
    }

    /// Attach a custom router to the microkernel server
    ///
    /// This method allows attaching arbitrary axum Routers to the microkernel server,
    /// enabling extension of the server with custom endpoints while maintaining
    /// microkernel architecture principles.
    ///
    /// # Arguments
    /// * `router` - The axum Router to attach
    ///
    /// # Examples
    /// ```rust
    /// use axum::{Router, routing::get, response::Html};
    /// use remote_mcp_kernel::microkernel::GitHubMicrokernelServer;
    /// 
    /// async fn hello() -> Html<&'static str> {
    ///     Html("<h1>Hello, World!</h1>")
    /// }
    /// 
    /// let custom_router = Router::new()
    ///     .route("/hello", get(hello));
    /// 
    /// let server: GitHubMicrokernelServer = GitHubMicrokernelServer::new()
    ///     .with_custom_router(custom_router);
    /// ```
    pub fn with_custom_router(mut self, router: Router) -> Self {
        let custom_router = CustomRouter {
            router,
            config: CustomRouterConfig::default(),
        };
        self.custom_routers.push(custom_router);
        self
    }

    /// Attach a custom router with configuration to the microkernel server
    ///
    /// This method provides more control over how the custom router is attached,
    /// allowing for custom configuration such as path prefixes and names.
    ///
    /// # Arguments
    /// * `router` - The axum Router to attach
    /// * `config` - Configuration for the custom router
    ///
    /// # Examples
    /// ```rust
    /// use axum::{Router, routing::get, response::Html};
    /// use remote_mcp_kernel::microkernel::{GitHubMicrokernelServer, CustomRouterConfig};
    /// 
    /// async fn api_status() -> Html<&'static str> {
    ///     Html("<h1>API Status: OK</h1>")
    /// }
    /// 
    /// let custom_router = Router::new()
    ///     .route("/status", get(api_status));
    /// 
    /// let config = CustomRouterConfig {
    ///     name: Some("API Router".to_string()),
    ///     path_prefix: Some("/api".to_string()),
    /// };
    /// 
    /// let server: GitHubMicrokernelServer = GitHubMicrokernelServer::new()
    ///     .with_custom_router_config(custom_router, config);
    /// ```
    pub fn with_custom_router_config(mut self, router: Router, config: CustomRouterConfig) -> Self {
        let custom_router = CustomRouter { router, config };
        self.custom_routers.push(custom_router);
        self
    }

    /// Attach multiple custom routers at once
    ///
    /// This method allows attaching multiple custom routers efficiently.
    ///
    /// # Arguments
    /// * `routers` - A vector of tuples containing (Router, CustomRouterConfig)
    ///
    /// # Examples
    /// ```rust
    /// use axum::{Router, routing::get, response::Html};
    /// use remote_mcp_kernel::microkernel::{GitHubMicrokernelServer, CustomRouterConfig};
    /// 
    /// async fn health() -> Html<&'static str> {
    ///     Html("<h1>Health: OK</h1>")
    /// }
    /// 
    /// async fn metrics() -> Html<&'static str> {
    ///     Html("<h1>Metrics: OK</h1>")
    /// }
    /// 
    /// let health_router = Router::new().route("/health", get(health));
    /// let metrics_router = Router::new().route("/metrics", get(metrics));
    /// 
    /// let routers = vec![
    ///     (health_router, CustomRouterConfig::default()),
    ///     (metrics_router, CustomRouterConfig {
    ///         name: Some("Metrics".to_string()),
    ///         path_prefix: Some("/monitoring".to_string()),
    ///     }),
    /// ];
    /// 
    /// let server: GitHubMicrokernelServer = GitHubMicrokernelServer::new()
    ///     .with_custom_routers(routers);
    /// ```
    pub fn with_custom_routers(mut self, routers: Vec<(Router, CustomRouterConfig)>) -> Self {
        for (router, config) in routers {
            let custom_router = CustomRouter { router, config };
            self.custom_routers.push(custom_router);
        }
        self
    }

    /// Build the composed router from all registered handlers
    ///
    /// This method demonstrates the microkernel composition principle
    /// where independent components are combined into a unified service.
    /// Now includes support for custom routers with path prefixes and configuration.
    pub fn build_router(self) -> AppResult<Router> {
        let mut router = Router::new();

        // Compose OAuth handler if present
        if let Some(oauth_provider) = self.oauth_provider {
            let oauth_router = oauth_provider.router();
            router = router.merge(oauth_router);
            tracing::debug!("✓ OAuth provider router composed");
        }

        // Compose Streamable HTTP handler if present
        if let Some(streamable_handler) = self.streamable_handler {
            let streamable_router = streamable_handler.router();
            router = router.merge(streamable_router);
            tracing::debug!("✓ Streamable HTTP handler router composed");
        }

        // Compose SSE handler if present
        if let Some(sse_handler) = self.sse_handler {
            let sse_router = sse_handler.router(self.sse_config)?;
            router = router.merge(sse_router);
            tracing::debug!("✓ SSE handler router composed");
        }

        // Compose custom routers
        for (index, custom_router) in self.custom_routers.into_iter().enumerate() {
            let CustomRouter { router: custom_router_inner, config } = custom_router;
            
            let final_router = if let Some(path_prefix) = config.path_prefix {
                // If a path prefix is specified, nest the router under that path
                Router::new().nest(&path_prefix, custom_router_inner)
            } else {
                // Otherwise, merge directly
                custom_router_inner
            };

            router = router.merge(final_router);
            
            // Log with custom name if provided, otherwise use index
            if let Some(name) = config.name {
                tracing::debug!("✓ Custom router '{}' composed", name);
            } else {
                tracing::debug!("✓ Custom router #{} composed", index + 1);
            }
        }

        Ok(router)
    }

    /// Start the microkernel server
    pub async fn serve(self, bind_address: SocketAddr) -> AppResult<()> {
        let router = self.build_router()?;

        let listener = tokio::net::TcpListener::bind(bind_address)
            .await
            .map_err(|e| crate::error::AppError::Internal(format!("Failed to bind: {}", e)))?;

        tracing::info!("🏛️  Microkernel server listening on {}", bind_address);

        axum::serve(listener, router)
            .await
            .map_err(|e| crate::error::AppError::Internal(format!("Server error: {}", e)))?;

        Ok(())
    }
}

impl<
    P: OAuthProviderTrait<S, DefaultClientManager<S>>,
    S: OAuthStorage + Clone + 'static,
    M: McpServerHandler,
> Default for MicrokernelServer<P, S, M>
{
    fn default() -> Self {
        Self::new()
    }
}

impl CustomRouterConfig {
    /// Create a new custom router configuration with a name
    pub fn with_name(name: impl Into<String>) -> Self {
        Self {
            name: Some(name.into()),
            path_prefix: None,
        }
    }

    /// Create a new custom router configuration with a path prefix
    pub fn with_prefix(path_prefix: impl Into<String>) -> Self {
        Self {
            name: None,
            path_prefix: Some(path_prefix.into()),
        }
    }

    /// Create a new custom router configuration with both name and path prefix
    pub fn with_name_and_prefix(name: impl Into<String>, path_prefix: impl Into<String>) -> Self {
        Self {
            name: Some(name.into()),
            path_prefix: Some(path_prefix.into()),
        }
    }
}