component_shape_mcp/

server.rs

use super::*;

/// A clonable collection of MCP tool definitions and in-process handlers.
///
/// Resources, prompts, server identity, and transport lifecycle remain owned
/// by [`McpServer`].
/// Output schemas are compiled at registration and shared by cloned registries.
/// Successful structured results are checked against the advertised schema on
/// direct, async, and protocol calls; handler error results bypass this check.
/// See [`tool_definition`] for dialect, format, and reference behavior.
#[derive(Clone, Default)]
pub struct McpToolRegistry {
    tools: BTreeMap<String, Arc<dyn ToolExecutor>>,
}

impl McpToolRegistry {
    /// Create an empty tool registry.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a synchronous MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_tool<Call>(
        &mut self,
        definition: ToolDefinition,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Call: Fn(McpToolCall) -> ToolCallResult + Send + Sync + 'static,
    {
        let name = definition.name.to_string();
        let output_validator =
            definitions::validate_tool_definition_with_output(&definition)?.map(Arc::new);
        if self.tools.contains_key(&name) {
            return Err(McpToolError::duplicate_tool(name));
        }

        self.tools.insert(
            name,
            Arc::new(RegisteredTool {
                definition,
                output_validator,
                call: Arc::new(move |arguments| Box::pin(std::future::ready(call(arguments)))),
            }),
        );
        Ok(())
    }

    /// Register a synchronous typed MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_typed_tool<Input, Call>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput,
        Call: Fn(Input) -> ToolCallResult + Send + Sync + 'static,
    {
        self.add_tool(definition.into_definition(), move |tool_call| {
            let input = match Input::from_tool_call(tool_call) {
                Ok(input) => input,
                Err(error) => return tool_error_result_for(error),
            };
            call(input)
        })
    }

    /// Register a typed handler with canonical Koruma domain validation.
    ///
    /// Strict decoding runs first, then [`validate_koruma`]. A failing value
    /// returns structured issues without invoking the handler. Successful
    /// results still pass the advertised output-schema checks.
    ///
    /// ```
    /// use koruma_collection::numeric::RangeValidation;
    ///
    /// #[derive(component_shape_mcp::McpToolInput, koruma::Koruma)]
    /// # #[mcp(crate = component_shape_mcp)]
    /// struct BatchArgs {
    ///     #[koruma(RangeValidation::<_>.min(1).max(5))]
    ///     count: u32,
    /// }
    /// # fn main() -> Result<(), component_shape_mcp::McpToolError> {
    /// let mut tools = component_shape_mcp::McpToolRegistry::new();
    /// let definition = component_shape_mcp::tool_definition_for_input::<BatchArgs>(
    ///     "batch", None, None, None,
    /// )?;
    /// tools.add_koruma_tool(definition, |args: BatchArgs| {
    ///     component_shape_mcp::tool_structured_result(
    ///         component_shape_mcp::serde_json::json!({ "count": args.count }),
    ///     )
    /// })?;
    /// assert_eq!(tools.call_tool("batch", Some(
    ///     component_shape_mcp::serde_json::json!({ "count": 0 }),
    /// )).is_error, Some(true));
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] for an invalid or duplicate tool definition.
    pub fn add_koruma_tool<Input, Call>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput + koruma_core::ValidateExt,
        Input::Error: koruma_core::ValidationIssues,
        Call: Fn(Input) -> ToolCallResult + Send + Sync + 'static,
    {
        self.add_typed_tool(definition, move |input| match validate_koruma(&input) {
            Ok(()) => call(input),
            Err(error) => tool_error_result_for(error),
        })
    }

    /// Register an async MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_tool_async<Call, Fut>(
        &mut self,
        definition: ToolDefinition,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Call: Fn(McpToolCall) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        let name = definition.name.to_string();
        let output_validator =
            definitions::validate_tool_definition_with_output(&definition)?.map(Arc::new);
        if self.tools.contains_key(&name) {
            return Err(McpToolError::duplicate_tool(name));
        }

        self.tools.insert(
            name,
            Arc::new(RegisteredTool {
                definition,
                output_validator,
                call: Arc::new(move |arguments| Box::pin(call(arguments))),
            }),
        );
        Ok(())
    }

    /// Register an async typed MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_typed_tool_async<Input, Call, Fut>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput,
        Call: Fn(Input) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.add_tool_async(definition.into_definition(), move |tool_call| {
            let input = Input::from_tool_call(tool_call);
            let future = match input {
                Ok(input) => call(input),
                Err(error) => {
                    return Box::pin(std::future::ready(tool_error_result_for(error)))
                        as ToolFuture;
                },
            };
            Box::pin(future) as ToolFuture
        })
    }

    /// Register an async typed handler with Koruma domain validation.
    ///
    /// Decoding and validation run before the handler creates its future.
    /// See [`Self::add_koruma_tool`] for issue and output-schema behavior.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] for an invalid or duplicate tool definition.
    pub fn add_koruma_tool_async<Input, Call, Fut>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput + koruma_core::ValidateExt,
        Input::Error: koruma_core::ValidationIssues,
        Call: Fn(Input) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.add_typed_tool_async(definition, move |input| match validate_koruma(&input) {
            Ok(()) => Box::pin(call(input)) as ToolFuture,
            Err(error) => Box::pin(std::future::ready(tool_error_result_for(error))) as ToolFuture,
        })
    }

    /// Return registered MCP tool definitions.
    #[must_use]
    pub fn list_tools(&self) -> Vec<ToolDefinition> {
        self.tools
            .values()
            .map(|executor| executor.definition())
            .collect()
    }

    /// Whether a tool name is already registered.
    #[must_use]
    pub fn contains_tool(&self, name: &str) -> bool {
        self.tools.contains_key(name)
    }

    /// Number of registered tools.
    #[must_use]
    pub fn tool_count(&self) -> usize {
        self.tools.len()
    }

    /// Calls a registered tool and converts validation or handler failures into
    /// a protocol-level tool result.
    #[must_use]
    pub fn call_tool(&self, name: &str, arguments: Option<Value>) -> ToolCallResult {
        let call = match McpToolCall::from_value(arguments) {
            Ok(call) => call,
            Err(error) => return tool_error_result_for(error),
        };
        let call = match self.resolve_tool(name, call) {
            Ok(resolved) => resolved,
            Err(error) => return tool_error_result_for(error),
        };
        block_on_tool_future(call)
    }

    /// Asynchronously calls a registered tool and converts validation or
    /// handler failures into a protocol-level tool result.
    pub async fn call_tool_async(&self, name: &str, arguments: Option<Value>) -> ToolCallResult {
        let call = match McpToolCall::from_value(arguments) {
            Ok(call) => call,
            Err(error) => return tool_error_result_for(error),
        };
        let call = match self.resolve_tool(name, call) {
            Ok(resolved) => resolved,
            Err(error) => return tool_error_result_for(error),
        };
        call.await
    }

    fn definition(&self, name: &str) -> Option<ToolDefinition> {
        self.tools.get(name).map(|executor| executor.definition())
    }

    fn resolve_tool(&self, name: &str, call: McpToolCall) -> Result<ToolFuture, McpToolError> {
        match self.tools.get(name) {
            Some(executor) => Ok(executor.call(call)),
            None => Err(McpToolError::UnknownTool {
                name: name.to_string(),
            }),
        }
    }
}

/// In-process MCP server that owns registered tools, resources, and prompts.
#[derive(Clone)]
pub struct McpServer {
    pub(crate) server_name: Cow<'static, str>,
    pub(crate) server_version: Cow<'static, str>,
    tools: McpToolRegistry,
    resources: BTreeMap<String, Arc<dyn ResourceReader>>,
    resource_templates: Vec<ResourceTemplate>,
    prompts: BTreeMap<String, Arc<dyn PromptExecutor>>,
}

impl McpServer {
    /// Create a dynamic MCP tool server with the advertised metadata.
    pub fn new(
        server_name: impl Into<Cow<'static, str>>,
        server_version: impl Into<Cow<'static, str>>,
    ) -> Self {
        Self {
            server_name: server_name.into(),
            server_version: server_version.into(),
            tools: McpToolRegistry::new(),
            resources: BTreeMap::new(),
            resource_templates: Vec::new(),
            prompts: BTreeMap::new(),
        }
    }

    /// Start building a dynamic MCP tool server with generated registrars.
    pub fn builder(
        server_name: impl Into<Cow<'static, str>>,
        server_version: impl Into<Cow<'static, str>>,
    ) -> McpServerBuilder {
        McpServerBuilder::new(server_name, server_version)
    }

    /// Create a dynamic MCP server from an existing tool registry.
    #[must_use]
    pub fn from_tool_registry(
        server_name: impl Into<Cow<'static, str>>,
        server_version: impl Into<Cow<'static, str>>,
        tools: McpToolRegistry,
    ) -> Self {
        Self {
            tools,
            ..Self::new(server_name, server_version)
        }
    }

    /// Borrow the shared tool registry.
    #[must_use]
    pub fn tool_registry(&self) -> &McpToolRegistry {
        &self.tools
    }

    /// Mutably borrow the shared tool registry.
    pub fn tool_registry_mut(&mut self) -> &mut McpToolRegistry {
        &mut self.tools
    }

    /// Consume the server and return its tool registry.
    #[must_use]
    pub fn into_tool_registry(self) -> McpToolRegistry {
        self.tools
    }

    /// Register a synchronous MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_tool<Call>(
        &mut self,
        definition: ToolDefinition,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Call: Fn(McpToolCall) -> ToolCallResult + Send + Sync + 'static,
    {
        self.tools.add_tool(definition, call)
    }

    /// Register a synchronous typed MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_typed_tool<Input, Call>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput,
        Call: Fn(Input) -> ToolCallResult + Send + Sync + 'static,
    {
        self.tools.add_typed_tool(definition, call)
    }

    /// Register a typed handler with Koruma validation before dispatch.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] for an invalid or duplicate tool definition.
    pub fn add_koruma_tool<Input, Call>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput + koruma_core::ValidateExt,
        Input::Error: koruma_core::ValidationIssues,
        Call: Fn(Input) -> ToolCallResult + Send + Sync + 'static,
    {
        self.tools.add_koruma_tool(definition, call)
    }

    /// Register an async MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_tool_async<Call, Fut>(
        &mut self,
        definition: ToolDefinition,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Call: Fn(McpToolCall) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.tools.add_tool_async(definition, call)
    }

    /// Register an async typed MCP tool handler.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a tool with the
    /// same name is already registered.
    pub fn add_typed_tool_async<Input, Call, Fut>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput,
        Call: Fn(Input) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.tools.add_typed_tool_async(definition, call)
    }

    /// Register an async typed handler with Koruma validation before dispatch.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] for an invalid or duplicate tool definition.
    pub fn add_koruma_tool_async<Input, Call, Fut>(
        &mut self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Result<(), McpToolError>
    where
        Input: McpToolInput + koruma_core::ValidateExt,
        Input::Error: koruma_core::ValidationIssues,
        Call: Fn(Input) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.tools.add_koruma_tool_async(definition, call)
    }

    /// Return registered MCP tool definitions.
    #[must_use]
    pub fn list_tools(&self) -> Vec<ToolDefinition> {
        self.tools.list_tools()
    }

    /// Whether a tool name is already registered.
    #[must_use]
    pub fn contains_tool(&self, name: &str) -> bool {
        self.tools.contains_tool(name)
    }

    /// Number of registered tools.
    #[must_use]
    pub fn tool_count(&self) -> usize {
        self.tools.tool_count()
    }

    /// Register a static MCP resource reader.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a resource with
    /// the same URI is already registered.
    pub fn add_resource<Read>(
        &mut self,
        definition: ResourceDefinition,
        read: Read,
    ) -> Result<(), McpToolError>
    where
        Read: Fn() -> ReadResourceResult + Send + Sync + 'static,
    {
        self.add_resource_async(definition, move || {
            let result = read();
            std::future::ready(Ok(result))
        })
    }

    /// Register an async MCP resource reader.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a resource with
    /// the same URI is already registered.
    pub fn add_resource_async<Read, Fut>(
        &mut self,
        definition: ResourceDefinition,
        read: Read,
    ) -> Result<(), McpToolError>
    where
        Read: Fn() -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ReadResourceResult, ErrorData>> + Send + 'static,
    {
        let uri = definition.uri.clone();
        validate_resource_definition(&definition)?;
        if self.resources.contains_key(&uri) {
            return Err(McpToolError::duplicate_resource(uri));
        }

        self.resources.insert(
            uri,
            Arc::new(RegisteredResource {
                definition,
                read: Arc::new(move || Box::pin(read())),
            }),
        );
        Ok(())
    }

    /// Register an MCP resource template advertised by `resources/templates/list`.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid.
    pub fn add_resource_template(
        &mut self,
        definition: ResourceTemplateDefinition,
    ) -> Result<(), McpToolError> {
        validate_resource_template(&definition)?;
        self.resource_templates.push(definition);
        Ok(())
    }

    /// Return registered MCP resource definitions.
    pub fn list_resources(&self) -> Vec<ResourceDefinition> {
        self.resources
            .values()
            .map(|resource| resource.definition())
            .collect()
    }

    /// Return registered MCP resource template definitions.
    pub fn list_resource_templates(&self) -> Vec<ResourceTemplateDefinition> {
        self.resource_templates.clone()
    }

    /// Whether a resource URI is already registered.
    pub fn contains_resource(&self, uri: &str) -> bool {
        self.resources.contains_key(uri)
    }

    /// Number of registered concrete resources.
    pub fn resource_count(&self) -> usize {
        self.resources.len()
    }

    /// Register a static MCP prompt.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a prompt with
    /// the same name is already registered.
    pub fn add_prompt<Get>(
        &mut self,
        definition: PromptDefinition,
        get: Get,
    ) -> Result<(), McpToolError>
    where
        Get: Fn(Option<JsonObject>) -> GetPromptResult + Send + Sync + 'static,
    {
        self.add_prompt_async(definition, move |arguments| {
            let result = get(arguments);
            std::future::ready(Ok(result))
        })
    }

    /// Register an async MCP prompt.
    ///
    /// # Errors
    ///
    /// Returns [`McpToolError`] when `definition` is invalid or a prompt with
    /// the same name is already registered.
    pub fn add_prompt_async<Get, Fut>(
        &mut self,
        definition: PromptDefinition,
        get: Get,
    ) -> Result<(), McpToolError>
    where
        Get: Fn(Option<JsonObject>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<GetPromptResult, ErrorData>> + Send + 'static,
    {
        let name = definition.name.clone();
        validate_prompt_definition(&definition)?;
        if self.prompts.contains_key(&name) {
            return Err(McpToolError::duplicate_prompt(name));
        }

        self.prompts.insert(
            name,
            Arc::new(RegisteredPrompt {
                definition,
                get: Arc::new(move |arguments| Box::pin(get(arguments))),
            }),
        );
        Ok(())
    }

    /// Return registered MCP prompt definitions.
    pub fn list_prompts(&self) -> Vec<PromptDefinition> {
        self.prompts
            .values()
            .map(|prompt| prompt.definition())
            .collect()
    }

    /// Whether a prompt name is already registered.
    pub fn contains_prompt(&self, name: &str) -> bool {
        self.prompts.contains_key(name)
    }

    /// Number of registered prompts.
    pub fn prompt_count(&self) -> usize {
        self.prompts.len()
    }

    /// Calls a registered tool and converts validation or handler failures into
    /// a protocol-level tool result.
    #[must_use]
    pub fn call_tool(&self, name: &str, arguments: Option<Value>) -> ToolCallResult {
        self.tools.call_tool(name, arguments)
    }

    /// Asynchronously calls a registered tool and converts validation or
    /// handler failures into a protocol-level tool result.
    pub async fn call_tool_async(&self, name: &str, arguments: Option<Value>) -> ToolCallResult {
        self.tools.call_tool_async(name, arguments).await
    }

    /// Serve this server over stdin/stdout using the MCP stdio transport.
    ///
    /// # Errors
    ///
    /// Returns an error when stdio serving fails or the service task fails.
    pub async fn serve_stdio(self) -> ServeStdioResult {
        let service = self.serve(stdio()).await?;
        service.waiting().await?;
        Ok(())
    }

    /// Serve this server over stdin/stdout on a new Tokio runtime.
    ///
    /// # Errors
    ///
    /// Returns an error when the Tokio runtime cannot be created, stdio serving
    /// fails, or the service task fails.
    pub fn serve_stdio_blocking(self) -> ServeStdioResult {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()?;
        runtime.block_on(self.serve_stdio())
    }

    fn result_meta(&self) -> MetaObject {
        let mut meta = MetaObject::default();
        meta.0.insert(
            "io.modelcontextprotocol/serverInfo".to_string(),
            json!({
                "name": self.server_name,
                "version": self.server_version,
            }),
        );
        meta
    }

    fn prepare_cacheable_result(
        &self,
        protocol_version: Option<&ProtocolVersion>,
        ttl_ms: &mut Option<u64>,
        cache_scope: &mut Option<CacheScope>,
        meta: &mut Option<MetaObject>,
    ) {
        if protocol_version.is_some_and(|version| version >= &ProtocolVersion::V_2026_07_28) {
            *ttl_ms = Some(0);
            *cache_scope = Some(CacheScope::Private);
        }
        *meta = Some(self.result_meta());
    }
}

/// Builder for composing generated MCP server registrars.
#[derive(Clone)]
pub struct McpServerBuilder {
    server: Result<McpServer, McpToolError>,
}

impl McpServerBuilder {
    /// Create a builder with advertised server metadata.
    pub fn new(
        server_name: impl Into<Cow<'static, str>>,
        server_version: impl Into<Cow<'static, str>>,
    ) -> Self {
        Self {
            server: Ok(McpServer::new(server_name, server_version)),
        }
    }

    /// Runs a registrar against the server being built.
    pub fn register<Register>(mut self, register: Register) -> Self
    where
        Register: FnOnce(&mut McpServer) -> Result<(), McpToolError>,
    {
        if let Ok(server) = self.server.as_mut()
            && let Err(error) = register(server)
        {
            self.server = Err(error);
        }
        self
    }

    /// Add a synchronous MCP tool to the server being built.
    pub fn tool<Call>(self, definition: ToolDefinition, call: Call) -> Self
    where
        Call: Fn(McpToolCall) -> ToolCallResult + Send + Sync + 'static,
    {
        self.register(move |server| server.add_tool(definition, call))
    }

    /// Add a synchronous typed MCP tool to the server being built.
    pub fn typed_tool<Input, Call>(self, definition: McpTypedTool<Input>, call: Call) -> Self
    where
        Input: McpToolInput,
        Call: Fn(Input) -> ToolCallResult + Send + Sync + 'static,
    {
        self.register(move |server| server.add_typed_tool(definition, call))
    }

    /// Add a typed tool with Koruma domain validation before its handler.
    pub fn koruma_tool<Input, Call>(self, definition: McpTypedTool<Input>, call: Call) -> Self
    where
        Input: McpToolInput + koruma_core::ValidateExt,
        Input::Error: koruma_core::ValidationIssues,
        Call: Fn(Input) -> ToolCallResult + Send + Sync + 'static,
    {
        self.register(move |server| server.add_koruma_tool(definition, call))
    }

    /// Add an async MCP tool to the server being built.
    pub fn tool_async<Call, Fut>(self, definition: ToolDefinition, call: Call) -> Self
    where
        Call: Fn(McpToolCall) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.register(move |server| server.add_tool_async(definition, call))
    }

    /// Add an async typed MCP tool to the server being built.
    pub fn typed_tool_async<Input, Call, Fut>(
        self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Self
    where
        Input: McpToolInput,
        Call: Fn(Input) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.register(move |server| server.add_typed_tool_async(definition, call))
    }

    /// Add an async typed tool with Koruma validation before creating its future.
    pub fn koruma_tool_async<Input, Call, Fut>(
        self,
        definition: McpTypedTool<Input>,
        call: Call,
    ) -> Self
    where
        Input: McpToolInput + koruma_core::ValidateExt,
        Input::Error: koruma_core::ValidationIssues,
        Call: Fn(Input) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ToolCallResult> + Send + 'static,
    {
        self.register(move |server| server.add_koruma_tool_async(definition, call))
    }

    /// Add a static MCP resource to the server being built.
    pub fn resource<Read>(self, definition: ResourceDefinition, read: Read) -> Self
    where
        Read: Fn() -> ReadResourceResult + Send + Sync + 'static,
    {
        self.register(move |server| server.add_resource(definition, read))
    }

    /// Add an async MCP resource to the server being built.
    pub fn resource_async<Read, Fut>(self, definition: ResourceDefinition, read: Read) -> Self
    where
        Read: Fn() -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ReadResourceResult, ErrorData>> + Send + 'static,
    {
        self.register(move |server| server.add_resource_async(definition, read))
    }

    /// Add an MCP resource template to the server being built.
    pub fn resource_template(self, definition: ResourceTemplateDefinition) -> Self {
        self.register(move |server| server.add_resource_template(definition))
    }

    /// Add a static MCP prompt to the server being built.
    pub fn prompt<Get>(self, definition: PromptDefinition, get: Get) -> Self
    where
        Get: Fn(Option<JsonObject>) -> GetPromptResult + Send + Sync + 'static,
    {
        self.register(move |server| server.add_prompt(definition, get))
    }

    /// Add an async MCP prompt to the server being built.
    pub fn prompt_async<Get, Fut>(self, definition: PromptDefinition, get: Get) -> Self
    where
        Get: Fn(Option<JsonObject>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<GetPromptResult, ErrorData>> + Send + 'static,
    {
        self.register(move |server| server.add_prompt_async(definition, get))
    }

    /// Finish the builder and return the composed server.
    ///
    /// # Errors
    ///
    /// Returns the first [`McpToolError`] produced by a builder registrar.
    pub fn build(self) -> Result<McpServer, McpToolError> {
        self.server
    }

    /// Build and serve this server over stdin/stdout using the MCP stdio transport.
    ///
    /// # Errors
    ///
    /// Returns an error when a builder registrar fails, stdio serving fails, or
    /// the service task fails.
    pub async fn serve_stdio(self) -> ServeStdioResult {
        self.build()?.serve_stdio().await
    }

    /// Build and serve this server over stdin/stdout on a new Tokio runtime.
    ///
    /// # Errors
    ///
    /// Returns an error when a builder registrar fails, the Tokio runtime cannot
    /// be created, stdio serving fails, or the service task fails.
    pub fn serve_stdio_blocking(self) -> ServeStdioResult {
        self.build()?.serve_stdio_blocking()
    }
}

impl ServerHandler for McpServer {
    fn get_info(&self) -> ServerConfig {
        let mut capabilities = ServerCapabilities::builder().enable_tools().build();
        capabilities.resources = (!self.resources.is_empty()
            || !self.resource_templates.is_empty())
        .then(Default::default);
        capabilities.prompts = (!self.prompts.is_empty()).then(Default::default);
        ServerConfig::new(capabilities)
            .with_protocol_version(ProtocolVersion::V_2026_07_28)
            .with_server_info(Implementation::new(
                self.server_name.clone(),
                self.server_version.clone(),
            ))
    }

    fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<ListToolsResult, ErrorData>> + MaybeSendFuture + '_ {
        let mut result = ListToolsResult::with_all_items(self.list_tools());
        self.prepare_cacheable_result(
            context.protocol_version().as_ref(),
            &mut result.ttl_ms,
            &mut result.cache_scope,
            &mut result.meta,
        );
        std::future::ready(Ok(result))
    }

    fn get_tool(&self, name: &str) -> Option<Tool> {
        self.tools.definition(name)
    }

    fn call_tool(
        &self,
        request: CallToolRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<CallToolResponse, ErrorData>> + MaybeSendFuture + '_ {
        let name = request.name.to_string();
        let call = McpToolCall::new(request.arguments.unwrap_or_default());
        let call = match self.tools.resolve_tool(&name, call) {
            Ok(call) => call,
            Err(error) => {
                let result = tool_error_result_for(error);
                Box::pin(std::future::ready(result))
            },
        };
        let result_meta = self.result_meta();
        async move {
            let mut result = call.await;
            result.meta = Some(result_meta);
            Ok(result.into())
        }
    }

    fn list_resources(
        &self,
        _request: Option<PaginatedRequestParams>,
        context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<ListResourcesResult, ErrorData>> + MaybeSendFuture + '_ {
        let mut result = ListResourcesResult::with_all_items(self.list_resources());
        self.prepare_cacheable_result(
            context.protocol_version().as_ref(),
            &mut result.ttl_ms,
            &mut result.cache_scope,
            &mut result.meta,
        );
        std::future::ready(Ok(result))
    }

    fn list_resource_templates(
        &self,
        _request: Option<PaginatedRequestParams>,
        context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<ListResourceTemplatesResult, ErrorData>> + MaybeSendFuture + '_
    {
        let mut result =
            ListResourceTemplatesResult::with_all_items(self.list_resource_templates());
        self.prepare_cacheable_result(
            context.protocol_version().as_ref(),
            &mut result.ttl_ms,
            &mut result.cache_scope,
            &mut result.meta,
        );
        std::future::ready(Ok(result))
    }

    fn read_resource(
        &self,
        request: ReadResourceRequestParams,
        context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<ReadResourceResponse, ErrorData>> + MaybeSendFuture + '_ {
        let uri = request.uri;
        let read = self.resources.get(&uri).map(|resource| resource.read());
        let protocol_version = context.protocol_version();
        let result_meta = self.result_meta();
        async move {
            match read {
                Some(read) => {
                    let mut result = read.await?;
                    if protocol_version
                        .as_ref()
                        .is_some_and(|version| version >= &ProtocolVersion::V_2026_07_28)
                    {
                        result.ttl_ms = Some(0);
                        result.cache_scope = Some(CacheScope::Private);
                    }
                    result.meta = Some(result_meta);
                    Ok(result.into())
                },
                None => Err(ErrorData::resource_not_found(
                    format!("resource `{uri}` not found"),
                    Some(McpToolError::unknown_resource(uri).to_structured_value()),
                )),
            }
        }
    }

    fn list_prompts(
        &self,
        _request: Option<PaginatedRequestParams>,
        context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<ListPromptsResult, ErrorData>> + MaybeSendFuture + '_ {
        let mut result = ListPromptsResult::with_all_items(self.list_prompts());
        self.prepare_cacheable_result(
            context.protocol_version().as_ref(),
            &mut result.ttl_ms,
            &mut result.cache_scope,
            &mut result.meta,
        );
        std::future::ready(Ok(result))
    }

    fn get_prompt(
        &self,
        request: GetPromptRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<GetPromptResponse, ErrorData>> + MaybeSendFuture + '_ {
        let name = request.name;
        let get = self
            .prompts
            .get(&name)
            .map(|prompt| prompt.get(request.arguments));
        let result_meta = self.result_meta();
        async move {
            match get {
                Some(get) => {
                    let mut result = get.await?;
                    result.meta = Some(result_meta);
                    Ok(result.into())
                },
                None => Err(ErrorData::invalid_params(
                    format!("prompt `{name}` not found"),
                    Some(McpToolError::unknown_prompt(name).to_structured_value()),
                )),
            }
        }
    }
}

trait ToolExecutor: Send + Sync {
    fn definition(&self) -> ToolDefinition;
    fn call(&self, call: McpToolCall) -> ToolFuture;
}

struct RegisteredTool {
    definition: ToolDefinition,
    output_validator: Option<Arc<jsonschema::Validator>>,
    call: Arc<dyn Fn(McpToolCall) -> ToolFuture + Send + Sync>,
}

impl ToolExecutor for RegisteredTool {
    fn definition(&self) -> ToolDefinition {
        self.definition.clone()
    }

    fn call(&self, call: McpToolCall) -> ToolFuture {
        let future = (self.call)(call);
        let Some(validator) = self.output_validator.clone() else {
            return future;
        };
        let name = self.definition.name.clone();
        Box::pin(async move { validate_tool_call_result(&name, &validator, future.await) })
    }
}

trait ResourceReader: Send + Sync {
    fn definition(&self) -> ResourceDefinition;
    fn read(&self) -> ResourceFuture;
}

struct RegisteredResource {
    definition: ResourceDefinition,
    read: Arc<dyn Fn() -> ResourceFuture + Send + Sync>,
}

impl ResourceReader for RegisteredResource {
    fn definition(&self) -> ResourceDefinition {
        self.definition.clone()
    }

    fn read(&self) -> ResourceFuture {
        (self.read)()
    }
}

trait PromptExecutor: Send + Sync {
    fn definition(&self) -> PromptDefinition;
    fn get(&self, arguments: Option<JsonObject>) -> PromptFuture;
}

struct RegisteredPrompt {
    definition: PromptDefinition,
    get: Arc<dyn Fn(Option<JsonObject>) -> PromptFuture + Send + Sync>,
}

impl PromptExecutor for RegisteredPrompt {
    fn definition(&self) -> PromptDefinition {
        self.definition.clone()
    }

    fn get(&self, arguments: Option<JsonObject>) -> PromptFuture {
        (self.get)(arguments)
    }
}

fn validate_tool_call_result(
    tool_name: &str,
    output_validator: &jsonschema::Validator,
    result: ToolCallResult,
) -> ToolCallResult {
    if result.is_error == Some(true) {
        return result;
    }

    let Some(structured_content) = result.structured_content.as_ref() else {
        return tool_error_result_for(McpToolError::invalid_tool_output(
            tool_name,
            "tool declares output_schema but returned no structured_content",
        ));
    };

    if let Err(error) = output_validator.validate(structured_content) {
        return tool_error_result_for(McpToolError::invalid_tool_output(
            tool_name,
            format!(
                "structured_content{}: {}",
                error.instance_path(),
                error.masked()
            ),
        ));
    }

    result
}

fn block_on_tool_future(future: ToolFuture) -> ToolCallResult {
    let join = std::thread::spawn(move || {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()?;
        Ok::<_, Box<dyn std::error::Error + Send + Sync>>(runtime.block_on(future))
    })
    .join();

    match join {
        Ok(Ok(result)) => result,
        Ok(Err(error)) => tool_error_result_for(McpToolError::handler(format!(
            "failed to run async tool handler: {error}"
        ))),
        Err(_) => {
            tool_error_result_for(McpToolError::handler("async tool handler runtime panicked"))
        },
    }
}