use std::fmt;
use std::pin::Pin;
use std::sync::Arc;
use rmcp::model::{
CallToolRequestParams, CallToolResponse, Implementation, ListToolsResult, PaginatedRequestParams,
ServerCapabilities, ServerConfig, Tool,
};
use rmcp::service::RequestContext;
use rmcp::{ErrorData, RoleServer, ServerHandler};
const NAME: &str = "dbmcp";
const VERSION: &str = env!("CARGO_PKG_VERSION");
const TITLE: &str = "Database MCP Server";
const HOMEPAGE: &str = env!("CARGO_PKG_HOMEPAGE");
#[derive(Clone)]
pub struct Server {
handler: Arc<dyn DynHandler>,
shutdown: Arc<dyn Shutdown>,
}
type HandlerFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, ErrorData>> + Send + 'a>>;
trait DynHandler: Send + Sync {
fn get_info(&self) -> ServerConfig;
fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> HandlerFuture<'_, CallToolResponse>;
fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> HandlerFuture<'_, ListToolsResult>;
fn get_tool(&self, name: &str) -> Option<Tool>;
}
impl<H: ServerHandler> DynHandler for H {
fn get_info(&self) -> ServerConfig {
ServerHandler::get_info(self)
}
fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> HandlerFuture<'_, CallToolResponse> {
Box::pin(ServerHandler::call_tool(self, request, context))
}
fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> HandlerFuture<'_, ListToolsResult> {
Box::pin(ServerHandler::list_tools(self, request, context))
}
fn get_tool(&self, name: &str) -> Option<Tool> {
ServerHandler::get_tool(self, name)
}
}
pub type ShutdownFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
pub trait Shutdown: Send + Sync {
fn shutdown(&self) -> ShutdownFuture<'_>;
}
impl fmt::Debug for Server {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Server").finish_non_exhaustive()
}
}
impl Server {
pub fn new<S>(handler: S) -> Self
where
S: ServerHandler + Shutdown,
{
let handler = Arc::new(handler);
Self {
handler: handler.clone(),
shutdown: handler,
}
}
pub async fn shutdown(&self) {
self.shutdown.shutdown().await;
}
}
impl ServerHandler for Server {
fn get_info(&self) -> ServerConfig {
self.handler.get_info()
}
async fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, ErrorData> {
self.handler.call_tool(request, context).await
}
async fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, ErrorData> {
self.handler.list_tools(request, context).await
}
fn get_tool(&self, name: &str) -> Option<Tool> {
self.handler.get_tool(name)
}
}
#[must_use]
pub fn server_info() -> ServerConfig {
let capabilities = ServerCapabilities::builder().enable_tools().build();
let server_info = Implementation::new(NAME, VERSION)
.with_title(TITLE)
.with_website_url(HOMEPAGE);
ServerConfig::new(capabilities).with_server_info(server_info)
}