Fix websocket ktor server implementation, add test and logs

Author: tiginamaria

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

@@ -158,14 +158,13 @@ public open class Client(
             serverVersion = result.serverInfo
 
             notification(InitializedNotification())
+        } catch (error: CancellationException) {
+            throw IllegalStateException("Error connecting to transport: ${error.message}")
         } catch (error: Throwable) {
+            logger.error(error) { "Failed to initialize client" }
             close()
-            if (error !is CancellationException) {
-                throw IllegalStateException("Error connecting to transport: ${error.message}")
-            }
 
             throw error
-
         }
     }
 

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

@@ -1,15 +1,17 @@
 package io.modelcontextprotocol.kotlin.sdk.client
 
-import io.ktor.client.HttpClient
-import io.ktor.client.plugins.websocket.webSocketSession
-import io.ktor.client.request.HttpRequestBuilder
-import io.ktor.client.request.header
-import io.ktor.http.HttpHeaders
-import io.ktor.websocket.WebSocketSession
+import io.github.oshai.kotlinlogging.KotlinLogging
+import io.ktor.client.*
+import io.ktor.client.plugins.websocket.*
+import io.ktor.client.request.*
+import io.ktor.http.*
+import io.ktor.websocket.*
 import io.modelcontextprotocol.kotlin.sdk.shared.MCP_SUBPROTOCOL
 import io.modelcontextprotocol.kotlin.sdk.shared.WebSocketMcpTransport
 import kotlin.properties.Delegates
 
+private val logger = KotlinLogging.logger {}
+
 /**
  * Client transport for WebSocket: this will connect to a server over the WebSocket protocol.
  */
@@ -21,6 +23,8 @@ public class WebSocketClientTransport(
     override var session: WebSocketSession by Delegates.notNull()
 
     override suspend fun initializeSession() {
+        logger.debug { "Websocket session initialization started..." }
+
         session = urlString?.let {
             client.webSocketSession(it) {
                 requestBuilder()
@@ -32,5 +36,7 @@ public class WebSocketClientTransport(
 
             header(HttpHeaders.SecWebSocketProtocol, MCP_SUBPROTOCOL)
         }
+
+        logger.debug { "Websocket session initialization finished" }
     }
 }

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

@@ -1,11 +1,15 @@
 package io.modelcontextprotocol.kotlin.sdk.client
 
+import io.github.oshai.kotlinlogging.KotlinLogging
 import io.ktor.client.HttpClient
 import io.ktor.client.request.HttpRequestBuilder
 import io.modelcontextprotocol.kotlin.sdk.Implementation
 import io.modelcontextprotocol.kotlin.sdk.LIB_VERSION
 import io.modelcontextprotocol.kotlin.sdk.shared.IMPLEMENTATION_NAME
 
+private val logger = KotlinLogging.logger {}
+
+
 /**
  * Returns a new WebSocket transport for the Model Context Protocol using the provided HttpClient.
  *
@@ -36,6 +40,8 @@ public suspend fun HttpClient.mcpWebSocket(
             version = LIB_VERSION
         )
     )
+    logger.debug { "Client started to connect to server" }
     client.connect(transport)
+    logger.debug { "Client finished to connect to server" }
     return client
 }

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

@@ -8,7 +8,6 @@ import io.ktor.server.response.respond
 import io.ktor.server.routing.Routing
 import io.ktor.server.routing.RoutingContext
 import io.ktor.server.routing.post
-import io.ktor.server.routing.route
 import io.ktor.server.routing.routing
 import io.ktor.server.sse.SSE
 import io.ktor.server.sse.ServerSSESession
@@ -18,39 +17,40 @@ import kotlinx.atomicfu.AtomicRef
 import kotlinx.atomicfu.atomic
 import kotlinx.atomicfu.update
 import kotlinx.collections.immutable.PersistentMap
-import kotlinx.collections.immutable.persistentMapOf
+import kotlinx.collections.immutable.toPersistentMap
 
 private val logger = KotlinLogging.logger {}
 
-@KtorDsl
-public fun Routing.mcp(path: String, block: () -> Server) {
-    route(path) {
-        mcp(block)
+internal class SseTransportManager(transports: Map<String, SseServerTransport> = emptyMap()) {
+    private val transports: AtomicRef<PersistentMap<String, SseServerTransport>> = atomic(transports.toPersistentMap())
+
+    fun getTransport(sessionId: String): SseServerTransport? = transports.value[sessionId]
+
+    fun addTransport(transport: SseServerTransport) {
+        transports.update { it.put(transport.sessionId, transport) }
+    }
+
+    fun removeTransport(sessionId: String) {
+        transports.update { it.remove(sessionId) }
     }
 }
 
-/**
- * Configures the Ktor Application to handle Model Context Protocol (MCP) over Server-Sent Events (SSE).
- */
+/*
+* Configures the Ktor Application to handle Model Context Protocol (MCP) over Server-Sent Events (SSE).
+*/
 @KtorDsl
 public fun Routing.mcp(block: () -> Server) {
-    val transports = atomic(persistentMapOf<String, SseServerTransport>())
+    val sseTransportManager = SseTransportManager()
 
     sse {
-        mcpSseEndpoint("", transports, block)
+        mcpSseEndpoint("", sseTransportManager, block)
     }
 
     post {
-        mcpPostEndpoint(transports)
+        mcpPostEndpoint(sseTransportManager)
     }
 }
 
-@Suppress("FunctionName")
-@Deprecated("Use mcp() instead", ReplaceWith("mcp(block)"), DeprecationLevel.WARNING)
-public fun Application.MCP(block: () -> Server) {
-    mcp(block)
-}
-
 @KtorDsl
 public fun Application.mcp(block: () -> Server) {
     install(SSE)
@@ -62,16 +62,16 @@ public fun Application.mcp(block: () -> Server) {
 
 internal suspend fun ServerSSESession.mcpSseEndpoint(
     postEndpoint: String,
-    transports: AtomicRef<PersistentMap<String, SseServerTransport>>,
+    sseTransportManager: SseTransportManager,
     block: () -> Server,
 ) {
-    val transport = mcpSseTransport(postEndpoint, transports)
+    val transport = mcpSseTransport(postEndpoint, sseTransportManager)
 
     val server = block()
 
     server.onClose {
         logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
-        transports.update { it.remove(transport.sessionId) }
+        sseTransportManager.removeTransport(transport.sessionId)
     }
 
     server.connectSession(transport)
@@ -81,17 +81,17 @@ internal suspend fun ServerSSESession.mcpSseEndpoint(
 
 internal fun ServerSSESession.mcpSseTransport(
     postEndpoint: String,
-    transports: AtomicRef<PersistentMap<String, SseServerTransport>>,
+    sseTransportManager: SseTransportManager,
 ): SseServerTransport {
     val transport = SseServerTransport(postEndpoint, this)
-    transports.update { it.put(transport.sessionId, transport) }
+    sseTransportManager.addTransport(transport)
     logger.info { "New SSE connection established and stored with sessionId: ${transport.sessionId}" }
 
     return transport
 }
 
 internal suspend fun RoutingContext.mcpPostEndpoint(
-    transports: AtomicRef<PersistentMap<String, SseServerTransport>>,
+    sseTransportManager: SseTransportManager,
 ) {
     val sessionId: String = call.request.queryParameters["sessionId"]
         ?: run {
@@ -101,7 +101,7 @@ internal suspend fun RoutingContext.mcpPostEndpoint(
 
     logger.debug { "Received message for sessionId: $sessionId" }
 
-    val transport = transports.value[sessionId]
+    val transport = sseTransportManager.getTransport(sessionId)
     if (transport == null) {
         logger.warn { "Session not found for sessionId: $sessionId" }
         call.respond(HttpStatusCode.NotFound, "Session not found")

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

@@ -191,7 +191,9 @@ public open class Server(
             }
         }
 
+        logger.debug { "Server session connecting to transport" }
         session.connect(transport)
+        logger.debug { "Server session successfully connected to transport" }
         sessions.update { it.add(session) }
 
         _onConnect()

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

@@ -92,7 +92,7 @@ public open class ServerSession(
      * Called when the server session is closing.
      */
     override fun onClose() {
-        logger.info { "Server connection closing" }
+        logger.debug { "Server connection closing" }
         _onClose()
     }
 
@@ -322,7 +322,7 @@ public open class ServerSession(
     }
 
     private suspend fun handleInitialize(request: InitializeRequest): InitializeResult {
-        logger.info { "Handling initialize request from client ${request.clientInfo}" }
+        logger.debug { "Handling initialization request from client" }
         clientCapabilities = request.capabilities
         clientVersion = request.clientInfo