refactor remove and add public tools/roots/prompts for server

Author: devcrocod

api/kotlin-sdk.api

@@ -2872,6 +2872,9 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server : io/modelcontextp
 	public static synthetic fun createMessage$default (Lio/modelcontextprotocol/kotlin/sdk/server/Server;Lio/modelcontextprotocol/kotlin/sdk/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public final fun getClientCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/ClientCapabilities;
 	public final fun getClientVersion ()Lio/modelcontextprotocol/kotlin/sdk/Implementation;
+	public final fun getPrompts ()Ljava/util/Map;
+	public final fun getResources ()Ljava/util/Map;
+	public final fun getTools ()Ljava/util/Map;
 	public final fun listRoots (Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun listRoots$default (Lio/modelcontextprotocol/kotlin/sdk/server/Server;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public fun onClose ()V

src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt

@@ -22,6 +22,7 @@ import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesRequest
 import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesResult
 import io.modelcontextprotocol.kotlin.sdk.ListResourcesRequest
 import io.modelcontextprotocol.kotlin.sdk.ListResourcesResult
+import io.modelcontextprotocol.kotlin.sdk.ListRootsRequest
 import io.modelcontextprotocol.kotlin.sdk.ListRootsResult
 import io.modelcontextprotocol.kotlin.sdk.ListToolsRequest
 import io.modelcontextprotocol.kotlin.sdk.ListToolsResult
@@ -44,7 +45,9 @@ import io.modelcontextprotocol.kotlin.sdk.shared.Transport
 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
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.JsonNull
 import kotlinx.serialization.json.JsonObject
@@ -105,7 +108,7 @@ public open class Client(
 
         // Internal handlers for roots
         if (capabilities.roots != null) {
-            setRequestHandler<ListToolsRequest>(Method.Defined.RootsList) { _, _ ->
+            setRequestHandler<ListRootsRequest>(Method.Defined.RootsList) { _, _ ->
                 handleListRoots()
             }
         }
@@ -544,12 +547,7 @@ public open class Client(
         }
         logger.info { "Removing ${uris.size} roots" }
 
-        val oldMap = roots.getAndUpdate { current ->
-            uris.fold(current) { map, uri ->
-                logger.debug { "Removing root: $uri" }
-                map.remove(uri)
-            }
-        }
+        val oldMap = roots.getAndUpdate { current -> current - uris.toPersistentSet() }
 
         val removedCount = uris.count { it in oldMap }
 

src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt

@@ -47,7 +47,9 @@ import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
 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
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.serialization.json.JsonObject
 
@@ -98,11 +100,11 @@ public open class Server(
     private val _tools = atomic(persistentMapOf<String, RegisteredTool>())
     private val _prompts = atomic(persistentMapOf<String, RegisteredPrompt>())
     private val _resources = atomic(persistentMapOf<String, RegisteredResource>())
-    private val tools: Map<String, RegisteredTool>
+    public val tools: Map<String, RegisteredTool>
         get() = _tools.value
-    private val prompts: Map<String, RegisteredPrompt>
+    public val prompts: Map<String, RegisteredPrompt>
         get() = _prompts.value
-    private val resources: Map<String, RegisteredResource>
+    public val resources: Map<String, RegisteredResource>
         get() = _resources.value
 
     init {
@@ -263,12 +265,7 @@ public open class Server(
         }
         logger.info { "Removing ${toolNames.size} tools" }
 
-        val oldMap = _tools.getAndUpdate { current ->
-            toolNames.fold(current) { map, name ->
-                logger.debug { "Removing tool: $name" }
-                map.remove(name)
-            }
-        }
+        val oldMap = _tools.getAndUpdate { current -> current - toolNames.toPersistentSet() }
 
         val removedCount = toolNames.count { it in oldMap }
         logger.info {
@@ -372,12 +369,7 @@ public open class Server(
         }
         logger.info { "Removing ${promptNames.size} prompts" }
 
-        val oldMap = _prompts.getAndUpdate { current ->
-            promptNames.fold(current) { map, name ->
-                logger.debug { "Removing prompt: $name" }
-                map.remove(name)
-            }
-        }
+        val oldMap = _prompts.getAndUpdate { current -> current - promptNames.toPersistentSet() }
 
         val removedCount = promptNames.count { it in oldMap }
 
@@ -477,12 +469,7 @@ public open class Server(
         }
         logger.info { "Removing ${uris.size} resources" }
 
-        val oldMap = _resources.getAndUpdate { current ->
-            uris.fold(current) { map, uri ->
-                logger.debug { "Removing resource: $uri" }
-                map.remove(uri)
-            }
-        }
+        val oldMap = _resources.getAndUpdate { current -> current - uris.toPersistentSet() }
 
         val removedCount = uris.count { it in oldMap }
 
@@ -611,7 +598,7 @@ public open class Server(
 
     private suspend fun handleCallTool(request: CallToolRequest): CallToolResult {
         logger.debug { "Handling tool call request for tool: ${request.name}" }
-        val tool = tools[request.name]
+        val tool = _tools.value[request.name]
             ?: run {
                 logger.error { "Tool not found: ${request.name}" }
                 throw IllegalArgumentException("Tool not found: ${request.name}")

src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt

@@ -239,7 +239,7 @@ public abstract class Protocol(
     private suspend fun onRequest(request: JSONRPCRequest) {
         LOGGER.trace { "Received request: ${request.method} (id: ${request.id})" }
 
-        val handler = _requestHandlers.value[request.method] ?: fallbackRequestHandler
+        val handler = requestHandlers[request.method] ?: fallbackRequestHandler
 
         if (handler === null) {
             LOGGER.trace { "No handler found for request: ${request.method}" }