feat: concurrent message dispatch across all transports

integration-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/BaseTransportTest.kt

@@ -1,12 +1,13 @@
 package io.modelcontextprotocol.kotlin.sdk.shared
 
 import io.kotest.assertions.nondeterministic.eventually
+import io.kotest.matchers.collections.shouldContainExactlyInAnyOrder
 import io.kotest.matchers.shouldBe
 import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification
 import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
 import io.modelcontextprotocol.kotlin.sdk.types.PingRequest
 import io.modelcontextprotocol.kotlin.sdk.types.toJSON
-import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.channels.Channel
 import kotlin.test.fail
 import kotlin.time.Duration.Companion.seconds
 
@@ -46,14 +47,10 @@ abstract class BaseTransportTest {
             InitializedNotification().toJSON(),
         )
 
-        val readMessages = mutableListOf<JSONRPCMessage>()
-        val finished = CompletableDeferred<Unit>()
+        val chan = Channel<JSONRPCMessage>(messages.size)
 
         transport.onMessage { message ->
-            readMessages.add(message)
-            if (message == messages.last()) {
-                finished.complete(Unit)
-            }
+            chan.send(message)
         }
 
         transport.start()
@@ -62,9 +59,13 @@ abstract class BaseTransportTest {
             transport.send(message)
         }
 
-        finished.await()
+        val readMessages = buildList {
+            repeat(messages.size) {
+                add(chan.receive())
+            }
+        }
 
-        messages shouldBe readMessages
+        readMessages shouldContainExactlyInAnyOrder messages
 
         transport.close()
     }

integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/kotlin/AbstractResourceIntegrationTest.kt

@@ -21,6 +21,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withContext
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.assertThrows
 import java.util.concurrent.CopyOnWriteArrayList
@@ -215,7 +216,7 @@ abstract class AbstractResourceIntegrationTest : KotlinTestBase() {
         val invalidUri = "test://nonexistent.txt"
 
         val exception = assertThrows<McpException> {
-            runBlocking {
+            withContext(Dispatchers.Default) {
                 client.readResource(ReadResourceRequest(ReadResourceRequestParams(uri = invalidUri)))
             }
         }

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

@@ -80,7 +80,7 @@ public class SseClientTransport(
             reconnectionTime = reconnectionTime,
             block = requestBuilder,
         )
-        scope = CoroutineScope(session.coroutineContext + SupervisorJob())
+        scope = CoroutineScope(session.coroutineContext + SupervisorJob(session.coroutineContext[Job]))
 
         job = scope.launch(CoroutineName("SseMcpClientTransport.connect#${hashCode()}")) {
             collectMessages()
@@ -163,17 +163,31 @@ public class SseClientTransport(
         }
     }
 
-    private suspend fun handleMessage(data: String) {
+    private fun handleMessage(data: String) {
         try {
             val message = McpJson.decodeFromString<JSONRPCMessage>(data)
-            _onMessage(message)
+            launchMessageHandler(message)
         } catch (e: SerializationException) {
             _onError(e)
         }
     }
 
+    private fun launchMessageHandler(message: JSONRPCMessage) {
+        scope.launch(CoroutineName("SseMcpClientTransport.message#${hashCode()}")) {
+            try {
+                _onMessage(message)
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Throwable) {
+                logger.error(e) { "Error processing message" }
+                _onError(e)
+            }
+        }
+    }
+
     override suspend fun closeResources() {
         withContext(NonCancellable) {
+            invokeOnCloseCallback()
             job?.cancel()
             try {
                 if (::session.isInitialized) session.cancel()

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

@@ -172,7 +172,7 @@ public class StdioClientTransport @JvmOverloads public constructor(
                     .collect { event ->
                         when (event) {
                             is Event.JsonRpc -> {
-                                handleJSONRPCMessage(event.message)
+                                launchMessageHandler(event.message)
                             }
 
                             is Event.StderrEvent -> {
@@ -274,14 +274,16 @@ public class StdioClientTransport @JvmOverloads public constructor(
         }
     }
 
-    private suspend fun handleJSONRPCMessage(msg: JSONRPCMessage) {
-        try {
-            _onMessage.invoke(msg)
-        } catch (e: CancellationException) {
-            throw e
-        } catch (e: Throwable) {
-            logger.error(e) { "Error processing message." }
-            runCatching { _onError.invoke(e) }
+    private fun launchMessageHandler(message: JSONRPCMessage) {
+        scope.launch(CoroutineName("StdioClientTransport.message#${hashCode()}")) {
+            try {
+                _onMessage.invoke(message)
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Throwable) {
+                logger.error(e) { "Error processing message." }
+                runCatching { _onError.invoke(e) }
+            }
         }
     }
 

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

@@ -37,12 +37,13 @@ import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import kotlin.math.pow
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.milliseconds
@@ -191,7 +192,7 @@ public class StreamableHttpClientTransport(
         when (response.contentType()?.withoutParameters()) {
             ContentType.Application.Json -> response.bodyAsText().takeIf { it.isNotEmpty() }?.let { json ->
                 runCatching { McpJson.decodeFromString<JSONRPCMessage>(json) }
-                    .onSuccess { _onMessage(it) }
+                    .onSuccess { launchMessageHandler(it) }
                     .onFailure {
                         _onError(it)
                         throw it
@@ -241,8 +242,11 @@ public class StreamableHttpClientTransport(
 
     override suspend fun closeResources() {
         logger.debug { "Client transport closing." }
-        sseJob?.cancelAndJoin()
-        scope.cancel()
+        withContext(NonCancellable) {
+            invokeOnCloseCallback()
+            sseJob?.cancel()
+            scope.coroutineContext[Job]?.cancelAndJoin()
+        }
     }
 
     /**
@@ -415,9 +419,9 @@ public class StreamableHttpClientTransport(
                                 .onSuccess { msg ->
                                     if (msg is JSONRPCResponse) receivedResponse = true
                                     if (replayMessageId != null && msg is JSONRPCResponse) {
-                                        _onMessage(msg.copy(id = replayMessageId))
+                                        launchMessageHandler(msg.copy(id = replayMessageId))
                                     } else {
-                                        _onMessage(msg)
+                                        launchMessageHandler(msg)
                                     }
                                 }
                                 .onFailure(_onError)
@@ -450,7 +454,7 @@ public class StreamableHttpClientTransport(
         var id: String? = null
         var eventName: String? = null
 
-        suspend fun dispatch(id: String?, eventName: String?, data: String) {
+        fun dispatch(id: String?, eventName: String?, data: String) {
             id?.let {
                 localLastEventId = it
                 hasPrimingEvent = true
@@ -464,9 +468,9 @@ public class StreamableHttpClientTransport(
                     .onSuccess { msg ->
                         if (msg is JSONRPCResponse) receivedResponse = true
                         if (replayMessageId != null && msg is JSONRPCResponse) {
-                            _onMessage(msg.copy(id = replayMessageId))
+                            launchMessageHandler(msg.copy(id = replayMessageId))
                         } else {
-                            _onMessage(msg)
+                            launchMessageHandler(msg)
                         }
                     }
                     .onFailure {
@@ -517,4 +521,16 @@ public class StreamableHttpClientTransport(
         }
         return SseStreamResult(hasPrimingEvent, receivedResponse, localLastEventId, localServerRetryDelay)
     }
+    private fun launchMessageHandler(message: JSONRPCMessage) {
+        scope.launch(CoroutineName("StreamableHttpTransport.message#${hashCode()}")) {
+            try {
+                _onMessage(message)
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Throwable) {
+                logger.error(e) { "Error processing message" }
+                _onError(e)
+            }
+        }
+    }
 }

kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/streamable/http/StreamableHttpClientTransportTest.kt

@@ -2,6 +2,7 @@ package io.modelcontextprotocol.kotlin.sdk.client.streamable.http
 
 import io.kotest.assertions.fail
 import io.kotest.assertions.nondeterministic.eventually
+import io.kotest.matchers.collections.shouldContainExactlyInAnyOrder
 import io.kotest.matchers.collections.shouldHaveSize
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.types.shouldBeInstanceOf
@@ -711,11 +712,10 @@ class StreamableHttpClientTransportTest {
 
         receivedMessages shouldHaveSize 2
 
-        val firstNotification = receivedMessages[0] as JSONRPCNotification
-        firstNotification.method shouldBe "notifications/progress"
-
-        val secondNotification = receivedMessages[1] as JSONRPCNotification
-        secondNotification.method shouldBe "notifications/tools/list_changed"
+        receivedMessages
+            .filterIsInstance<JSONRPCNotification>()
+            .map { it.method }
+            .shouldContainExactlyInAnyOrder("notifications/progress", "notifications/tools/list_changed")
 
         transport.close()
     }

kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt

@@ -10,10 +10,14 @@ import io.modelcontextprotocol.kotlin.sdk.types.McpJson
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.InternalCoroutinesApi
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.ClosedReceiveChannelException
 import kotlinx.coroutines.job
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import kotlin.concurrent.atomics.AtomicBoolean
 import kotlin.concurrent.atomics.ExperimentalAtomicApi
 import kotlin.coroutines.cancellation.CancellationException
@@ -30,7 +34,7 @@ private val logger = KotlinLogging.logger {}
 @OptIn(ExperimentalAtomicApi::class)
 public abstract class WebSocketMcpTransport : AbstractTransport() {
     private val scope by lazy {
-        CoroutineScope(session.coroutineContext + SupervisorJob())
+        CoroutineScope(session.coroutineContext + SupervisorJob(session.coroutineContext.job))
     }
 
     private val initialized: AtomicBoolean = AtomicBoolean(false)
@@ -73,8 +77,8 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
                 }
 
                 try {
-                    val message = McpJson.decodeFromString<JSONRPCMessage>(message.readText())
-                    _onMessage.invoke(message)
+                    val parsedMessage = McpJson.decodeFromString<JSONRPCMessage>(message.readText())
+                    launchMessageHandler(parsedMessage)
                 } catch (e: CancellationException) {
                     throw e
                 } catch (e: Exception) {
@@ -86,6 +90,11 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
 
         @OptIn(InternalCoroutinesApi::class)
         session.coroutineContext.job.invokeOnCompletion {
+            // Cancel the scope when the session completes. For normal session
+            // completion the SupervisorJob parent does not auto-cancel children;
+            // for error/cancellation the propagation already cancels the scope
+            // job, making this cancel a no-op.
+            scope.cancel()
             if (it != null) {
                 _onError.invoke(it)
             } else {
@@ -94,6 +103,19 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
         }
     }
 
+    private fun launchMessageHandler(message: JSONRPCMessage) {
+        scope.launch(CoroutineName("WebSocketMcpTransport.message#${hashCode()}")) {
+            try {
+                _onMessage.invoke(message)
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Throwable) {
+                logger.error(e) { "Error processing message" }
+                _onError.invoke(e)
+            }
+        }
+    }
+
     override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
         logger.debug { "Sending message" }
         if (!initialized.load()) {
@@ -109,7 +131,10 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
         }
 
         logger.debug { "Closing websocket session" }
-        session.close()
-        session.coroutineContext.job.join()
+        withContext(NonCancellable) {
+            invokeOnCloseCallback()
+            session.close()
+            scope.coroutineContext.job.cancelAndJoin()
+        }
     }
 }