use crate::handler::{PromptHandler, ResourceHandler, ServerHandler, TaskHandler, ToolHandler};
use mcpkit_core::capability::ServerCapabilities;
#[derive(Debug, Clone, Copy, Default)]
pub struct NotRegistered;
#[derive(Debug)]
pub struct Registered<T>(pub T);
pub struct ServerBuilder<H, Tools, Resources, Prompts, Tasks> {
handler: H,
tools: Tools,
resources: Resources,
prompts: Prompts,
tasks: Tasks,
capabilities: ServerCapabilities,
}
impl<H: ServerHandler>
ServerBuilder<H, NotRegistered, NotRegistered, NotRegistered, NotRegistered>
{
#[must_use]
pub fn new(handler: H) -> Self {
let capabilities = handler.capabilities();
Self {
handler,
tools: NotRegistered,
resources: NotRegistered,
prompts: NotRegistered,
tasks: NotRegistered,
capabilities,
}
}
}
impl<H, T, R, P, K> ServerBuilder<H, T, R, P, K>
where
H: ServerHandler,
{
#[must_use]
pub fn capabilities(mut self, caps: ServerCapabilities) -> Self {
self.capabilities = caps;
self
}
#[must_use]
pub const fn get_capabilities(&self) -> &ServerCapabilities {
&self.capabilities
}
}
impl<H, R, P, K> ServerBuilder<H, NotRegistered, R, P, K>
where
H: ServerHandler,
{
#[must_use]
pub fn with_tools<TH: ToolHandler>(
self,
tools: TH,
) -> ServerBuilder<H, Registered<TH>, R, P, K> {
ServerBuilder {
handler: self.handler,
tools: Registered(tools),
resources: self.resources,
prompts: self.prompts,
tasks: self.tasks,
capabilities: self.capabilities.with_tools(),
}
}
}
impl<H, T, P, K> ServerBuilder<H, T, NotRegistered, P, K>
where
H: ServerHandler,
{
#[must_use]
pub fn with_resources<RH: ResourceHandler>(
self,
resources: RH,
) -> ServerBuilder<H, T, Registered<RH>, P, K> {
ServerBuilder {
handler: self.handler,
tools: self.tools,
resources: Registered(resources),
prompts: self.prompts,
tasks: self.tasks,
capabilities: self.capabilities.with_resources(),
}
}
}
impl<H, T, R, K> ServerBuilder<H, T, R, NotRegistered, K>
where
H: ServerHandler,
{
#[must_use]
pub fn with_prompts<PH: PromptHandler>(
self,
prompts: PH,
) -> ServerBuilder<H, T, R, Registered<PH>, K> {
ServerBuilder {
handler: self.handler,
tools: self.tools,
resources: self.resources,
prompts: Registered(prompts),
tasks: self.tasks,
capabilities: self.capabilities.with_prompts(),
}
}
}
impl<H, T, R, P> ServerBuilder<H, T, R, P, NotRegistered>
where
H: ServerHandler,
{
#[must_use]
pub fn with_tasks<KH: TaskHandler>(
self,
tasks: KH,
) -> ServerBuilder<H, T, R, P, Registered<KH>> {
ServerBuilder {
handler: self.handler,
tools: self.tools,
resources: self.resources,
prompts: self.prompts,
tasks: Registered(tasks),
capabilities: self.capabilities,
}
}
}
impl<H, T, R, P, K> ServerBuilder<H, T, R, P, K>
where
H: ServerHandler + Send + Sync + 'static,
T: Send + Sync + 'static,
R: Send + Sync + 'static,
P: Send + Sync + 'static,
K: Send + Sync + 'static,
{
#[must_use]
pub fn build(self) -> Server<H, T, R, P, K> {
Server {
handler: self.handler,
tools: self.tools,
resources: self.resources,
prompts: self.prompts,
tasks: self.tasks,
capabilities: self.capabilities,
}
}
}
pub struct Server<H, T, R, P, K> {
handler: H,
tools: T,
resources: R,
prompts: P,
tasks: K,
capabilities: ServerCapabilities,
}
impl<H, T, R, P, K> Server<H, T, R, P, K>
where
H: ServerHandler,
{
#[must_use]
pub const fn capabilities(&self) -> &ServerCapabilities {
&self.capabilities
}
#[must_use]
pub const fn handler(&self) -> &H {
&self.handler
}
#[must_use]
pub fn server_info(&self) -> mcpkit_core::capability::ServerInfo {
self.handler.server_info()
}
}
impl<H, TH, R, P, K> Server<H, Registered<TH>, R, P, K>
where
H: ServerHandler,
TH: ToolHandler,
{
#[must_use]
pub const fn tool_handler(&self) -> &TH {
&self.tools.0
}
}
impl<H, T, RH, P, K> Server<H, T, Registered<RH>, P, K>
where
H: ServerHandler,
RH: ResourceHandler,
{
#[must_use]
pub const fn resource_handler(&self) -> &RH {
&self.resources.0
}
}
impl<H, T, R, PH, K> Server<H, T, R, Registered<PH>, K>
where
H: ServerHandler,
PH: PromptHandler,
{
#[must_use]
pub const fn prompt_handler(&self) -> &PH {
&self.prompts.0
}
}
impl<H, T, R, P, KH> Server<H, T, R, P, Registered<KH>>
where
H: ServerHandler,
KH: TaskHandler,
{
#[must_use]
pub const fn task_handler(&self) -> &KH {
&self.tasks.0
}
}
pub type FullServer<H, TH, RH, PH, KH> =
Server<H, Registered<TH>, Registered<RH>, Registered<PH>, Registered<KH>>;
pub type MinimalServer<H> = Server<H, NotRegistered, NotRegistered, NotRegistered, NotRegistered>;
#[cfg(test)]
mod tests {
use super::*;
use crate::context::Context;
use crate::handler::ToolHandler;
use mcpkit_core::capability::ServerInfo;
use mcpkit_core::error::McpError;
use mcpkit_core::types::{Tool, ToolOutput};
use serde_json::Value;
struct TestHandler;
impl ServerHandler for TestHandler {
fn server_info(&self) -> ServerInfo {
ServerInfo::new("test", "1.0.0")
}
fn capabilities(&self) -> ServerCapabilities {
ServerCapabilities::default()
}
}
struct TestToolHandler;
impl ToolHandler for TestToolHandler {
async fn list_tools(&self, _ctx: &Context<'_>) -> Result<Vec<Tool>, McpError> {
Ok(vec![])
}
async fn call_tool(
&self,
_name: &str,
_args: Value,
_ctx: &Context<'_>,
) -> Result<ToolOutput, McpError> {
Ok(ToolOutput::text("test"))
}
}
#[test]
fn test_server_builder_minimal() {
let server = ServerBuilder::new(TestHandler).build();
assert_eq!(server.server_info().name, "test");
assert_eq!(server.server_info().version, "1.0.0");
}
#[test]
fn test_server_builder_with_tools() {
let server = ServerBuilder::new(TestHandler)
.with_tools(TestToolHandler)
.build();
assert!(server.capabilities().has_tools());
let _tool_handler: &TestToolHandler = server.tool_handler();
}
#[test]
fn test_typestate_prevents_double_registration() {
let _server = ServerBuilder::new(TestHandler)
.with_tools(TestToolHandler)
.build();
}
#[test]
fn test_builder_order_independence() {
let server1 = ServerBuilder::new(TestHandler)
.with_tools(TestToolHandler)
.build();
let _server2: Server<
TestHandler,
Registered<TestToolHandler>,
NotRegistered,
NotRegistered,
NotRegistered,
> = ServerBuilder::new(TestHandler)
.with_tools(TestToolHandler)
.build();
assert!(server1.capabilities().has_tools());
}
}