Move features(tools, prompts, resources) management to the FeatureRegistry

kotlin-sdk-server/api/kotlin-sdk-server.api

@@ -5,40 +5,43 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/KtorServerKt {
 	public static final fun mcp (Lio/ktor/server/routing/Routing;Lkotlin/jvm/functions/Function1;)V
 }
 
-public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt {
+public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt : io/modelcontextprotocol/kotlin/sdk/server/Feature {
 	public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/Prompt;Lkotlin/jvm/functions/Function2;)V
 	public final fun component1 ()Lio/modelcontextprotocol/kotlin/sdk/Prompt;
 	public final fun component2 ()Lkotlin/jvm/functions/Function2;
 	public final fun copy (Lio/modelcontextprotocol/kotlin/sdk/Prompt;Lkotlin/jvm/functions/Function2;)Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt;
 	public static synthetic fun copy$default (Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt;Lio/modelcontextprotocol/kotlin/sdk/Prompt;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt;
 	public fun equals (Ljava/lang/Object;)Z
+	public fun getKey ()Ljava/lang/String;
 	public final fun getMessageProvider ()Lkotlin/jvm/functions/Function2;
 	public final fun getPrompt ()Lio/modelcontextprotocol/kotlin/sdk/Prompt;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
 }
 
-public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredResource {
+public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredResource : io/modelcontextprotocol/kotlin/sdk/server/Feature {
 	public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/Resource;Lkotlin/jvm/functions/Function2;)V
 	public final fun component1 ()Lio/modelcontextprotocol/kotlin/sdk/Resource;
 	public final fun component2 ()Lkotlin/jvm/functions/Function2;
 	public final fun copy (Lio/modelcontextprotocol/kotlin/sdk/Resource;Lkotlin/jvm/functions/Function2;)Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredResource;
 	public static synthetic fun copy$default (Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredResource;Lio/modelcontextprotocol/kotlin/sdk/Resource;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredResource;
 	public fun equals (Ljava/lang/Object;)Z
+	public fun getKey ()Ljava/lang/String;
 	public final fun getReadHandler ()Lkotlin/jvm/functions/Function2;
 	public final fun getResource ()Lio/modelcontextprotocol/kotlin/sdk/Resource;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
 }
 
-public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredTool {
+public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredTool : io/modelcontextprotocol/kotlin/sdk/server/Feature {
 	public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/Tool;Lkotlin/jvm/functions/Function2;)V
 	public final fun component1 ()Lio/modelcontextprotocol/kotlin/sdk/Tool;
 	public final fun component2 ()Lkotlin/jvm/functions/Function2;
 	public final fun copy (Lio/modelcontextprotocol/kotlin/sdk/Tool;Lkotlin/jvm/functions/Function2;)Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredTool;
 	public static synthetic fun copy$default (Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredTool;Lio/modelcontextprotocol/kotlin/sdk/Tool;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/server/RegisteredTool;
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getHandler ()Lkotlin/jvm/functions/Function2;
+	public fun getKey ()Ljava/lang/String;
 	public final fun getTool ()Lio/modelcontextprotocol/kotlin/sdk/Tool;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;

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

@@ -0,0 +1,120 @@
+package io.modelcontextprotocol.kotlin.sdk.server
+
+import io.github.oshai.kotlinlogging.KotlinLogging
+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
+
+/**
+ * A generic registry for managing features of a specified type. This class provides thread-safe
+ * operations for adding, removing, and retrieving features from the registry.
+ *
+ * @param T The type of the feature, constrained to implement the [Feature] interface.
+ * @param featureType A string description of the type of feature being managed.
+ *                    Used primarily for logging purposes.
+ */
+internal class FeatureRegistry<T : Feature>(private val featureType: String) {
+
+    private val logger = KotlinLogging.logger(name = "FeatureRegistry[$featureType]")
+
+    /**
+     * Atomic variable used to maintain a thread-safe registry of features.
+     * Stores a persistent map where each feature is identified by a unique key.
+     */
+    private val registry = atomic(persistentMapOf<FeatureKey, T>())
+
+    /**
+     * Provides a snapshot of all features currently registered in the registry.
+     * Keys represent unique identifiers for each feature, and values represent the associated features.
+     */
+    internal val values: Map<FeatureKey, T>
+        get() = registry.value
+
+    /**
+     * Adds the specified feature to the registry.
+     *
+     * @param feature The feature to be added to the registry.
+     */
+    internal fun add(feature: T) {
+        logger.info { "Adding $featureType: \"${feature.key}\"" }
+        registry.update { current -> current.put(feature.key, feature) }
+        logger.info { "Added $featureType: \"${feature.key}\"" }
+    }
+
+    /**
+     * Adds the given list of features to the registry. Each feature is mapped by its key
+     * and added to the current registry.
+     *
+     * @param features The list of features to add to the registry.
+     */
+    internal fun addAll(features: List<T>) {
+        logger.info { "Adding ${featureType}s: ${features.size}" }
+        registry.update { current -> current.putAll(features.associateBy { it.key }) }
+        logger.info { "Added ${featureType}s: ${features.size}" }
+    }
+
+    /**
+     * Removes the feature associated with the given key from the registry.
+     *
+     * @param key The key of the feature to be removed.
+     * @return `true` if the feature was successfully removed, or `false` if the feature was not found.
+     */
+    internal fun remove(key: FeatureKey): Boolean {
+        logger.info { "Removing $featureType: \"$key\"" }
+        val oldMap = registry.getAndUpdate { current -> current.remove(key) }
+
+        val removed = key in oldMap
+        logger.info {
+            if (removed) {
+                "Removed $featureType: \"$key\""
+            } else {
+                "$featureType not found: \"$key\""
+            }
+        }
+        return key in oldMap
+    }
+
+    /**
+     * Removes the features associated with the given keys from the registry.
+     *
+     * @param keys The list of keys whose associated features are to be removed.
+     * @return The number of features that were successfully removed.
+     */
+    internal fun removeAll(keys: List<FeatureKey>): Int {
+        logger.info { "Removing ${featureType}s: ${keys.size}" }
+        val oldMap = registry.getAndUpdate { current -> current - keys.toPersistentSet() }
+
+        val removedCount = keys.count { it in oldMap }
+        logger.info {
+            if (removedCount > 0) {
+                "Removed ${featureType}s: $removedCount"
+            } else {
+                "No $featureType were removed"
+            }
+        }
+
+        return removedCount
+    }
+
+    /**
+     * Retrieves the feature associated with the given key from the registry.
+     *
+     * @param key The key of the feature to retrieve.
+     * @return The feature associated with the given key, or `null` if no such feature exists in the registry.
+     */
+    internal fun get(key: FeatureKey): T? {
+        logger.info { "Getting $featureType: \"$key\"" }
+        val feature = registry.value[key]
+        logger.info {
+            if (feature != null) {
+                "Got $featureType: \"$key\""
+            } else {
+                "$featureType not found: \"$key\""
+            }
+        }
+        return feature
+    }
+}

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

@@ -0,0 +1,56 @@
+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
+
+internal typealias FeatureKey = String
+
+/**
+ * Represents a feature with an associated unique key.
+ */
+internal interface Feature {
+    val key: FeatureKey
+}
+
+/**
+ * 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) : Feature {
+    override val key: String = tool.name
+}
+
+/**
+ * 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,
+) : Feature {
+    override val key: String = prompt.name
+}
+
+/**
+ * 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,
+) : Feature {
+    override val key: String = resource.uri
+}