Introduce server session

Author: tiginamaria

kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt

@@ -1,5 +1,6 @@
 package io.modelcontextprotocol.kotlin.sdk.client
 
+import io.github.oshai.kotlinlogging.KotlinLogging
 import io.ktor.client.HttpClient
 import io.ktor.client.plugins.sse.ClientSSESession
 import io.ktor.client.plugins.sse.sseSession
@@ -46,6 +47,8 @@ public class SseClientTransport(
     private val reconnectionTime: Duration? = null,
     private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
 ) : AbstractTransport() {
+    private val logger = KotlinLogging.logger {}
+
     private val initialized: AtomicBoolean = AtomicBoolean(false)
     private val endpoint = CompletableDeferred<String>()
 
@@ -68,6 +71,7 @@ public class SseClientTransport(
         check(initialized.compareAndSet(expectedValue = false, newValue = true)) {
             "SSEClientTransport already started! If using Client class, note that connect() calls start() automatically."
         }
+        logger.info { "Starting SseClientTransport..." }
 
         try {
             session = urlString?.let {
@@ -111,6 +115,8 @@ public class SseClientTransport(
                 val text = response.bodyAsText()
                 error("Error POSTing to endpoint (HTTP ${response.status}): $text")
             }
+
+            logger.debug { "Client successfully sent message via SSE $endpoint" }
         } catch (e: Throwable) {
             _onError(e)
             throw e
@@ -157,6 +163,7 @@ public class SseClientTransport(
             val path = if (eventData.startsWith("/")) eventData.substring(1) else eventData
             val endpointUrl = Url("$baseUrl/$path")
             endpoint.complete(endpointUrl.toString())
+            logger.debug { "Client connected to endpoint: $endpointUrl" }
         } catch (e: Throwable) {
             _onError(e)
             endpoint.completeExceptionally(e)

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

@@ -13,8 +13,12 @@ import io.ktor.server.routing.routing
 import io.ktor.server.sse.SSE
 import io.ktor.server.sse.ServerSSESession
 import io.ktor.server.sse.sse
-import io.ktor.util.collections.ConcurrentMap
 import io.ktor.utils.io.KtorDsl
+import kotlinx.atomicfu.AtomicRef
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.update
+import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.persistentMapOf
 
 private val logger = KotlinLogging.logger {}
 
@@ -30,7 +34,7 @@ public fun Routing.mcp(path: String, block: () -> Server) {
  */
 @KtorDsl
 public fun Routing.mcp(block: () -> Server) {
-    val transports = ConcurrentMap<String, SseServerTransport>()
+    val transports = atomic(persistentMapOf<String, SseServerTransport>())
 
     sse {
         mcpSseEndpoint("", transports, block)
@@ -49,24 +53,16 @@ public fun Application.MCP(block: () -> Server) {
 
 @KtorDsl
 public fun Application.mcp(block: () -> Server) {
-    val transports = ConcurrentMap<String, SseServerTransport>()
-
     install(SSE)
 
     routing {
-        sse("/sse") {
-            mcpSseEndpoint("/message", transports, block)
-        }
-
-        post("/message") {
-            mcpPostEndpoint(transports)
-        }
+        mcp(block)
     }
 }
 
-private suspend fun ServerSSESession.mcpSseEndpoint(
+internal suspend fun ServerSSESession.mcpSseEndpoint(
     postEndpoint: String,
-    transports: ConcurrentMap<String, SseServerTransport>,
+    transports: AtomicRef<PersistentMap<String, SseServerTransport>>,
     block: () -> Server,
 ) {
     val transport = mcpSseTransport(postEndpoint, transports)
@@ -75,27 +71,27 @@ private suspend fun ServerSSESession.mcpSseEndpoint(
 
     server.onClose {
         logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
-        transports.remove(transport.sessionId)
+        transports.update { it.remove(transport.sessionId) }
     }
 
-    server.connect(transport)
+    server.connectSession(transport)
+
     logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" }
 }
 
 internal fun ServerSSESession.mcpSseTransport(
     postEndpoint: String,
-    transports: ConcurrentMap<String, SseServerTransport>,
+    transports: AtomicRef<PersistentMap<String, SseServerTransport>>,
 ): SseServerTransport {
     val transport = SseServerTransport(postEndpoint, this)
-    transports[transport.sessionId] = transport
-
+    transports.update { it.put(transport.sessionId, transport) }
     logger.info { "New SSE connection established and stored with sessionId: ${transport.sessionId}" }
 
     return transport
 }
 
 internal suspend fun RoutingContext.mcpPostEndpoint(
-    transports: ConcurrentMap<String, SseServerTransport>,
+    transports: AtomicRef<PersistentMap<String, SseServerTransport>>,
 ) {
     val sessionId: String = call.request.queryParameters["sessionId"]
         ?: run {
@@ -105,7 +101,7 @@ internal suspend fun RoutingContext.mcpPostEndpoint(
 
     logger.debug { "Received message for sessionId: $sessionId" }
 
-    val transport = transports[sessionId]
+    val transport = transports.value[sessionId]
     if (transport == null) {
         logger.warn { "Session not found for sessionId: $sessionId" }
         call.respond(HttpStatusCode.NotFound, "Session not found")