use axum::Router;
use std::net::SocketAddr;
use crate::{
error::AppResult,
handlers::{McpServerHandler, SseHandler, SseHandlerConfig, StreamableHttpHandler},
};
use oauth_provider_rs::{DefaultClientManager, OAuthProvider, OAuthProviderTrait, OAuthStorage};
#[derive(Debug, Clone)]
pub struct CustomRouterConfig {
pub name: Option<String>,
pub path_prefix: Option<String>,
}
impl Default for CustomRouterConfig {
fn default() -> Self {
Self {
name: None,
path_prefix: None,
}
}
}
pub struct CustomRouter {
router: Router,
config: CustomRouterConfig,
}
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>
{
pub fn new() -> Self {
Self {
oauth_provider: None,
streamable_handler: None,
sse_handler: None,
sse_config: SseHandlerConfig::default(),
custom_routers: Vec::new(),
}
}
pub fn with_oauth_provider(mut self, oauth_provider: OAuthProvider<P, S>) -> Self {
self.oauth_provider = Some(oauth_provider);
self
}
pub fn with_streamable_handler(mut self, streamable_handler: StreamableHttpHandler<M>) -> Self {
self.streamable_handler = Some(streamable_handler);
self
}
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
}
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
}
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
}
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
}
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
}
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
}
pub fn build_router(self) -> AppResult<Router> {
let mut router = Router::new();
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");
}
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");
}
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");
}
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 {
Router::new().nest(&path_prefix, custom_router_inner)
} else {
custom_router_inner
};
router = router.merge(final_router);
if let Some(name) = config.name {
tracing::debug!("✓ Custom router '{}' composed", name);
} else {
tracing::debug!("✓ Custom router #{} composed", index + 1);
}
}
Ok(router)
}
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 {
pub fn with_name(name: impl Into<String>) -> Self {
Self {
name: Some(name.into()),
path_prefix: None,
}
}
pub fn with_prefix(path_prefix: impl Into<String>) -> Self {
Self {
name: None,
path_prefix: Some(path_prefix.into()),
}
}
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()),
}
}
}