Move session management logic to ServerSessionRegistry.kt

Rewrite server session redirection to receivers

Author: tiginamaria

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

@@ -75,8 +75,6 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server {
 	public final fun getPrompts ()Ljava/util/Map;
 	public final fun getResources ()Ljava/util/Map;
 	protected final fun getServerInfo ()Lio/modelcontextprotocol/kotlin/sdk/Implementation;
-	public final fun getSession (Ljava/lang/String;)Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;
-	public final fun getSessionOrThrow (Ljava/lang/String;)Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;
 	public final fun getSessions ()Ljava/util/Map;
 	public final fun getTools ()Ljava/util/Map;
 	public final fun listRoots (Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

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

@@ -36,9 +36,6 @@ import io.modelcontextprotocol.kotlin.sdk.ToolAnnotations
 import io.modelcontextprotocol.kotlin.sdk.shared.ProtocolOptions
 import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
 import io.modelcontextprotocol.kotlin.sdk.shared.Transport
-import kotlinx.atomicfu.atomic
-import kotlinx.atomicfu.update
-import kotlinx.collections.immutable.persistentMapOf
 import kotlinx.coroutines.CancellationException
 import kotlinx.serialization.json.JsonObject
 
@@ -88,24 +85,13 @@ public open class Server(
         block: Server.() -> Unit = {},
     ) : this(serverInfo, options, { instructions }, block)
 
-    private val sessionRegistry = atomic(persistentMapOf<String, ServerSession>())
+    private val sessionRegistry = ServerSessionRegistry()
 
     /**
-     * Returns a read-only view of the current server sessions.
+     * Provides a snapshot of all sessions currently registered in the server
      */
-    public val sessions: Map<String, ServerSession>
-        get() = sessionRegistry.value
-
-    /**
-     * Gets a server session by its ID.
-     */
-    public fun getSession(sessionId: String): ServerSession? = sessions[sessionId]
-
-    /**
-     * Gets a server session by its ID or throws an exception if the session doesn't exist.
-     */
-    public fun getSessionOrThrow(sessionId: String): ServerSession =
-        sessions[sessionId] ?: throw IllegalArgumentException("Session not found: $sessionId")
+    public val sessions: Map<ServerSessionKey, ServerSession>
+        get() = sessionRegistry.sessions
 
     @Suppress("ktlint:standard:backing-property-naming")
     private var _onInitialized: (() -> Unit) = {}
@@ -200,12 +186,12 @@ public open class Server(
         // Register cleanup handler to remove session from list when it closes
         session.onClose {
             logger.debug { "Removing closed session from active sessions list" }
-            sessionRegistry.update { sessions -> sessions.remove(session.sessionId) }
+            sessionRegistry.removeSession(session.sessionId)
         }
         logger.debug { "Server session connecting to transport" }
         session.connect(transport)
         logger.debug { "Server session successfully connected to transport" }
-        sessionRegistry.update { sessions -> sessions.put(session.sessionId, session) }
+        sessionRegistry.addSession(session)
 
         _onConnect()
         return session
@@ -571,9 +557,8 @@ public open class Server(
      * Triggers [ServerSession.ping] request for session by provided [sessionId].
      * @param sessionId The session ID to ping
      */
-    public suspend fun ping(sessionId: String): EmptyRequestResult {
-        val session = getSessionOrThrow(sessionId)
-        return session.ping()
+    public suspend fun ping(sessionId: String): EmptyRequestResult = with(sessionRegistry.getSession(sessionId)) {
+        ping()
     }
 
     /**
@@ -589,9 +574,8 @@ public open class Server(
         sessionId: String,
         params: CreateMessageRequest,
         options: RequestOptions? = null,
-    ): CreateMessageResult {
-        val session = getSessionOrThrow(sessionId)
-        return session.request(params, options)
+    ): CreateMessageResult = with(sessionRegistry.getSession(sessionId)) {
+        request(params, options)
     }
 
     /**
@@ -607,9 +591,8 @@ public open class Server(
         sessionId: String,
         params: JsonObject = EmptyJsonObject,
         options: RequestOptions? = null,
-    ): ListRootsResult {
-        val session = getSessionOrThrow(sessionId)
-        return session.listRoots(params, options)
+    ): ListRootsResult = with(sessionRegistry.getSession(sessionId)) {
+        listRoots(params, options)
     }
 
     /**
@@ -627,9 +610,8 @@ public open class Server(
         message: String,
         requestedSchema: RequestedSchema,
         options: RequestOptions? = null,
-    ): CreateElicitationResult {
-        val session = getSessionOrThrow(sessionId)
-        return session.createElicitation(message, requestedSchema, options)
+    ): CreateElicitationResult = with(sessionRegistry.getSession(sessionId)) {
+        createElicitation(message, requestedSchema, options)
     }
 
     /**
@@ -639,8 +621,9 @@ public open class Server(
      * @param notification The logging message notification.
      */
     public suspend fun sendLoggingMessage(sessionId: String, notification: LoggingMessageNotification) {
-        val session = getSessionOrThrow(sessionId)
-        session.sendLoggingMessage(notification)
+        with(sessionRegistry.getSession(sessionId)) {
+            sendLoggingMessage(notification)
+        }
     }
 
     /**
@@ -650,8 +633,9 @@ public open class Server(
      * @param notification Details of the updated resource.
      */
     public suspend fun sendResourceUpdated(sessionId: String, notification: ResourceUpdatedNotification) {
-        val session = getSessionOrThrow(sessionId)
-        session.sendResourceUpdated(notification)
+        with(sessionRegistry.getSession(sessionId)) {
+            sendResourceUpdated(notification)
+        }
     }
 
     /**
@@ -660,8 +644,9 @@ public open class Server(
      * @param sessionId The session ID to send the resource list changed notification to.
      */
     public suspend fun sendResourceListChanged(sessionId: String) {
-        val session = getSessionOrThrow(sessionId)
-        session.sendResourceListChanged()
+        with(sessionRegistry.getSession(sessionId)) {
+            sendResourceListChanged()
+        }
     }
 
     /**
@@ -670,8 +655,9 @@ public open class Server(
      * @param sessionId The session ID to send the tool list changed notification to.
      */
     public suspend fun sendToolListChanged(sessionId: String) {
-        val session = getSessionOrThrow(sessionId)
-        session.sendToolListChanged()
+        with(sessionRegistry.getSession(sessionId)) {
+            sendToolListChanged()
+        }
     }
 
     /**
@@ -680,8 +666,9 @@ public open class Server(
      * @param sessionId The session ID to send the prompt list changed notification to.
      */
     public suspend fun sendPromptListChanged(sessionId: String) {
-        val session = getSessionOrThrow(sessionId)
-        session.sendPromptListChanged()
+        with(sessionRegistry.getSession(sessionId)) {
+            sendPromptListChanged()
+        }
     }
     // End the ServerSession redirection section
 }

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

@@ -0,0 +1,60 @@
+package io.modelcontextprotocol.kotlin.sdk.server
+
+import io.github.oshai.kotlinlogging.KotlinLogging
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.update
+import kotlinx.collections.immutable.persistentMapOf
+
+internal typealias ServerSessionKey = String
+
+/**
+ * Represents a registry for managing server sessions.
+ */
+internal class ServerSessionRegistry {
+
+    private val logger = KotlinLogging.logger {}
+
+    /**
+     * Atomic variable used to maintain a thread-safe registry of sessions.
+     * Stores a persistent map where each session is identified by its unique key.
+     */
+    private val registry = atomic(persistentMapOf<String, ServerSession>())
+
+    /**
+     * Returns a read-only view of the current server sessions.
+     */
+    internal val sessions: Map<ServerSessionKey, ServerSession>
+        get() = registry.value
+
+    /**
+     * Returns a server session by its ID.
+     * @param sessionId The ID of the session to retrieve.
+     * @throws IllegalArgumentException If the session doesn't exist.
+     */
+    internal fun getSession(sessionId: ServerSessionKey): ServerSession =
+        sessions[sessionId] ?: throw IllegalArgumentException("Session not found: $sessionId")
+
+    /**
+     * Returns a server session by its ID, or null if it doesn't exist.
+     * @param sessionId The ID of the session to retrieve.
+     */
+    internal fun getSessionOrNull(sessionId: ServerSessionKey): ServerSession? = sessions[sessionId]
+
+    /**
+     * Registers a server session.
+     * @param session The session to register.
+     */
+    internal fun addSession(session: ServerSession) {
+        logger.info { "Adding session: ${session.sessionId}" }
+        registry.update { sessions -> sessions.put(session.sessionId, session) }
+    }
+
+    /**
+     * Removes a server session by its ID.
+     * @param sessionId The ID of the session to remove.
+     */
+    internal fun removeSession(sessionId: ServerSessionKey) {
+        logger.info { "Removing session: $sessionId" }
+        registry.update { sessions -> sessions.remove(sessionId) }
+    }
+}