use anyhow::Result;
use kodegen_config::CATEGORY_SEQUENTIAL_THINKING;
use kodegen_server_http::{ServerBuilder, Managers, RouterSet, ShutdownHook, register_tool};
use rmcp::handler::server::router::{prompt::PromptRouter, tool::ToolRouter};
use std::sync::Arc;
struct SequentialThinkingWrapper(Arc<kodegen_tools_sequential_thinking::SequentialThinkingTool>);
impl ShutdownHook for SequentialThinkingWrapper {
fn shutdown(&self) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
let tool = self.0.clone();
Box::pin(async move {
tool.shutdown().await
.map_err(|e| anyhow::anyhow!("Failed to shutdown sequential thinking tool: {e}"))
})
}
}
#[tokio::main]
async fn main() -> Result<()> {
ServerBuilder::new()
.category(CATEGORY_SEQUENTIAL_THINKING)
.register_tools(|| async {
let mut tool_router = ToolRouter::new();
let mut prompt_router = PromptRouter::new();
let managers = Managers::new();
let tool = kodegen_tools_sequential_thinking::SequentialThinkingTool::new();
let tool_arc = Arc::new(tool.clone());
tool_arc.clone().start_cleanup_task();
managers.register(SequentialThinkingWrapper(tool_arc)).await;
(tool_router, prompt_router) = register_tool(
tool_router,
prompt_router,
tool,
);
Ok(RouterSet::new(tool_router, prompt_router, managers))
})
.run()
.await
}