Move features management to the FeatureRegistry

Author: tiginamaria

kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureRegistry.kt

@@ -0,0 +1,220 @@
+package io.modelcontextprotocol.kotlin.sdk.server
+
+import io.github.oshai.kotlinlogging.KotlinLogging
+import io.modelcontextprotocol.kotlin.sdk.CallToolRequest
+import io.modelcontextprotocol.kotlin.sdk.CallToolResult
+import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
+import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
+import io.modelcontextprotocol.kotlin.sdk.Prompt
+import io.modelcontextprotocol.kotlin.sdk.ReadResourceRequest
+import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
+import io.modelcontextprotocol.kotlin.sdk.Resource
+import io.modelcontextprotocol.kotlin.sdk.Tool
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.getAndUpdate
+import kotlinx.atomicfu.update
+import kotlinx.collections.immutable.minus
+import kotlinx.collections.immutable.persistentMapOf
+import kotlinx.collections.immutable.toPersistentSet
+
+/**
+ * FeaturesRegistry is an internal class responsible for managing the registration and retrieval of
+ * tools [RegisteredTool], prompts [RegisteredPrompt], and resources [RegisteredResource].
+ * It provides mechanisms to add, remove, and fetch these entities through a thread-safe approach using atomic updates.
+ */
+internal class FeatureRegistry {
+    private val logger = KotlinLogging.logger {}
+
+    private val toolsRegistry = atomic(persistentMapOf<String, RegisteredTool>())
+    private val promptsRegistry = atomic(persistentMapOf<String, RegisteredPrompt>())
+    private val resourcesRegistry = atomic(persistentMapOf<String, RegisteredResource>())
+
+    /**
+     * A read-only map of currently registered tools, where the key is the tool's unique name,
+     * and the value is the corresponding [RegisteredTool] object containing its definition
+     * and handler.
+     */
+    internal val tools: Map<String, RegisteredTool>
+        get() = toolsRegistry.value
+
+    /**
+     * A read-only map of currently registered prompts, where the key is the prompt's unique name,
+     * and the value is the corresponding [RegisteredPrompt] object containing its definition
+     * and handler.
+     */
+    internal val prompts: Map<String, RegisteredPrompt>
+        get() = promptsRegistry.value
+
+    /**
+     * A read-only map of currently registered resources, where the key is the resource's URI,
+     * and the value is the corresponding [RegisteredResource] object containing its definition
+     * and handler.
+     */
+    internal val resources: Map<String, RegisteredResource>
+        get() = resourcesRegistry.value
+
+    //region Tools
+
+    /**
+     * Registers a single tool.
+     *
+     * @param tool A [io.modelcontextprotocol.kotlin.sdk.Tool] object describing the tool.
+     * @param handler A suspend function that handles executing the tool when called by the client.
+     */
+    internal fun addTool(tool: Tool, handler: suspend (CallToolRequest) -> CallToolResult) {
+        logger.info { "Registering tool: ${tool.name}" }
+        toolsRegistry.update { current -> current.put(tool.name, RegisteredTool(tool, handler)) }
+    }
+
+    /**
+     * Registers multiple tools at once.
+     *
+     * @param tools A list of [RegisteredTool] objects representing the tools to register.
+     */
+    internal fun addTools(tools: List<RegisteredTool>) {
+        logger.info { "Registering ${tools.size} tools" }
+        toolsRegistry.update { current -> current.putAll(tools.associateBy { it.tool.name }) }
+    }
+
+    /**
+     * Removes a tool from the registry by its name.
+     *
+     * @param toolName The name of the tool to be removed.
+     * @return True if the tool was successfully removed, false if the tool was not found.
+     */
+    internal fun removeTool(toolName: String): Boolean {
+        logger.info { "Unregistering tool: $toolName" }
+        val oldMap = toolsRegistry.getAndUpdate { current -> current.remove(toolName) }
+
+        return toolName in oldMap
+    }
+
+    /**
+     * Removes all tools from the registry that match the provided list of tool names.
+     *
+     * @param toolNames A list of names of the tools to be removed from the registry.
+     * @return The number of tools that were successfully removed.
+     */
+    internal fun removeTools(toolNames: List<String>): Int {
+        logger.info { "Unregistering ${toolNames.size} tools" }
+        val oldMap = toolsRegistry.getAndUpdate { current -> current - toolNames.toPersistentSet() }
+
+        val removedCount = toolNames.count { it in oldMap }
+
+        return removedCount
+    }
+
+    internal fun getTool(toolName: String): RegisteredTool? = toolsRegistry.value[toolName]
+
+    //endregion Tools
+
+    //region Prompts
+
+    /**
+     * Registers a single prompt.
+     *
+     * @param prompt A [Prompt] object describing the prompt.
+     * @param promptProvider A suspend function that returns the prompt content when requested by the client.
+     */
+    internal fun addPrompt(prompt: Prompt, promptProvider: suspend (GetPromptRequest) -> GetPromptResult) {
+        logger.info { "Registering prompt: ${prompt.name}" }
+        promptsRegistry.update { current -> current.put(prompt.name, RegisteredPrompt(prompt, promptProvider)) }
+    }
+
+    /**
+     * Registers multiple prompts at once.
+     *
+     * @param prompts A list of [RegisteredPrompt] objects representing the prompts to register.
+     */
+    internal fun addPrompts(prompts: List<RegisteredPrompt>) {
+        logger.info { "Registering ${prompts.size} prompts" }
+        promptsRegistry.update { current -> current.putAll(prompts.associateBy { it.prompt.name }) }
+    }
+
+    /**
+     * Removes a prompt from the registry by its name.
+     *
+     * @param promptName The name of the prompt to be removed.
+     * @return True if the prompt was successfully removed, false if the prompt was not found.
+     */
+    internal fun removePrompt(promptName: String): Boolean {
+        logger.info { "Unregistering prompt: $promptName" }
+        val oldMap = promptsRegistry.getAndUpdate { current -> current.remove(promptName) }
+
+        return promptName in oldMap
+    }
+
+    /**
+     * Removes all prompts from the registry that match the provided list of prompt names.
+     *
+     * @param promptNames A list of names of the prompts to be removed from the registry.
+     * @return The number of prompts that were successfully removed.
+     */
+    internal fun removePrompts(promptNames: List<String>): Int {
+        logger.info { "Unregistering ${promptNames.size} prompts" }
+        val oldMap = promptsRegistry.getAndUpdate { current -> current - promptNames.toPersistentSet() }
+
+        val removedCount = promptNames.count { it in oldMap }
+
+        return removedCount
+    }
+
+    internal fun getPrompt(promptName: String): RegisteredPrompt? = promptsRegistry.value[promptName]
+
+    //endregion Prompts
+
+    //region Resources
+
+    /**
+     * Registers a single resource.
+     *
+     * @param resource A [Resource] object describing the resource.
+     * @param readHandler A suspend function that handles reading the resource when requested by the client.
+     */
+    internal fun addResource(resource: Resource, readHandler: suspend (ReadResourceRequest) -> ReadResourceResult) {
+        logger.info { "Registering resource: ${resource.name} (${resource.uri})" }
+        resourcesRegistry.update { current -> current.put(resource.uri, RegisteredResource(resource, readHandler)) }
+    }
+
+    /**
+     * Registers multiple resources at once.
+     *
+     * @param resources A list of [RegisteredResource] objects representing the resources to register.
+     */
+    internal fun addResources(resources: List<RegisteredResource>) {
+        logger.info { "Registering ${resources.size} resources" }
+        resourcesRegistry.update { current -> current.putAll(resources.associateBy { it.resource.uri }) }
+    }
+
+    /**
+     * Removes a resource from the registry by its URI.
+     *
+     * @param resourceUri The URI of the resource to be removed.
+     * @return True if the resource was successfully removed, false if the resource was not found.
+     */
+    internal fun removeResource(resourceUri: String): Boolean {
+        logger.info { "Unregistering resource: $resourceUri" }
+        val oldMap = resourcesRegistry.getAndUpdate { current -> current.remove(resourceUri) }
+
+        return resourceUri in oldMap
+    }
+
+    /**
+     * Removes all resources from the registry that match the provided list of resource URIs.
+     *
+     * @param resourceUris A list of URIs of the resources to be removed from the registry.
+     * @return The number of resources that were successfully removed.
+     */
+    internal fun removeResources(resourceUris: List<String>): Int {
+        logger.info { "Unregistering ${resourceUris.size} resources" }
+        val oldMap = resourcesRegistry.getAndUpdate { current -> current - resourceUris.toPersistentSet() }
+
+        val removedCount = resourceUris.count { it in oldMap }
+
+        return removedCount
+    }
+
+    internal fun getResource(resourceUri: String): RegisteredResource? = resourcesRegistry.value[resourceUri]
+
+    //endregion Resources
+}

kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Features.kt

@@ -0,0 +1,41 @@
+package io.modelcontextprotocol.kotlin.sdk.server
+
+import io.modelcontextprotocol.kotlin.sdk.CallToolRequest
+import io.modelcontextprotocol.kotlin.sdk.CallToolResult
+import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
+import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
+import io.modelcontextprotocol.kotlin.sdk.Prompt
+import io.modelcontextprotocol.kotlin.sdk.ReadResourceRequest
+import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
+import io.modelcontextprotocol.kotlin.sdk.Resource
+import io.modelcontextprotocol.kotlin.sdk.Tool
+
+/**
+ * A wrapper class representing a registered tool on the server.
+ *
+ * @property tool The tool definition.
+ * @property handler A suspend function to handle the tool call requests.
+ */
+public data class RegisteredTool(val tool: Tool, val handler: suspend (CallToolRequest) -> CallToolResult)
+
+/**
+ * A wrapper class representing a registered prompt on the server.
+ *
+ * @property prompt The prompt definition.
+ * @property messageProvider A suspend function that returns the prompt content when requested by the client.
+ */
+public data class RegisteredPrompt(
+    val prompt: Prompt,
+    val messageProvider: suspend (GetPromptRequest) -> GetPromptResult,
+)
+
+/**
+ * A wrapper class representing a registered resource on the server.
+ *
+ * @property resource The resource definition.
+ * @property readHandler A suspend function to handle read requests for this resource.
+ */
+public data class RegisteredResource(
+    val resource: Resource,
+    val readHandler: suspend (ReadResourceRequest) -> ReadResourceResult,
+)