diff --git a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SseClientTransport.kt b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SseClientTransport.kt index 3c9d88b..a48db60 100644 --- a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SseClientTransport.kt +++ b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SseClientTransport.kt @@ -24,11 +24,12 @@ import kotlinx.coroutines.CoroutineName import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.Job +import kotlinx.coroutines.NonCancellable import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.cancel -import kotlinx.coroutines.cancelAndJoin import kotlinx.coroutines.ensureActive import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext import kotlinx.serialization.SerializationException import kotlin.concurrent.atomics.ExperimentalAtomicApi import kotlin.time.Duration @@ -171,13 +172,15 @@ public class SseClientTransport( } override suspend fun closeResources() { - job?.cancelAndJoin() - try { - if (::session.isInitialized) session.cancel() - if (::scope.isInitialized) scope.cancel() - endpoint.cancel() - } catch (e: Throwable) { - _onError(e) + withContext(NonCancellable) { + job?.cancel() + try { + if (::session.isInitialized) session.cancel() + if (::scope.isInitialized) scope.cancel() + endpoint.cancel() + } catch (e: Throwable) { + _onError(e) + } } } } diff --git a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport.kt b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport.kt index 12529cc..a61ea4b 100644 --- a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport.kt +++ b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport.kt @@ -21,6 +21,7 @@ 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 @@ -31,6 +32,7 @@ import kotlinx.coroutines.channels.consumeEach import kotlinx.coroutines.flow.channelFlow import kotlinx.coroutines.isActive import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext import kotlinx.coroutines.yield import kotlinx.io.Buffer import kotlinx.io.IOException @@ -241,8 +243,10 @@ public class StdioClientTransport @JvmOverloads public constructor( } override suspend fun closeResources() { - scope.stopProcessing("Closed") - scope.coroutineContext[Job]?.join() // Wait for all coroutines to complete + withContext(NonCancellable) { + scope.stopProcessing("Closed") + scope.coroutineContext[Job]?.join() // Wait for all coroutines to complete + } } private fun sendOutboundMessage(message: JSONRPCMessage, sink: Sink, mainScope: CoroutineScope) { @@ -264,6 +268,8 @@ 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) } diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt index 43c0ae5..bc3cf76 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt @@ -277,12 +277,15 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } try { handler(notification) + } catch (e: CancellationException) { + throw e } catch (cause: Throwable) { logger.error(cause) { "Error handling notification: ${notification.method}" } onError(cause) } } + @Suppress("ThrowsCount") private suspend fun onRequest(request: JSONRPCRequest) { logger.trace { "Received request: ${request.method} (id: ${request.id})" } @@ -300,6 +303,8 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio ), ), ) + } catch (e: CancellationException) { + throw e } catch (cause: Throwable) { logger.error(cause) { "Error sending method not found response" } onError(cause) diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt index 5f8bf9d..9fa1a0b 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt @@ -16,6 +16,7 @@ import kotlinx.coroutines.job import kotlinx.coroutines.launch import kotlin.concurrent.atomics.AtomicBoolean import kotlin.concurrent.atomics.ExperimentalAtomicApi +import kotlin.coroutines.cancellation.CancellationException /** WebSocket subprotocol identifier for MCP connections. */ public const val MCP_SUBPROTOCOL: String = "mcp" @@ -44,6 +45,7 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { */ protected abstract suspend fun initializeSession() + @Suppress("ThrowsCount") override suspend fun start() { logger.debug { "Starting websocket transport" } @@ -74,6 +76,8 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { try { val message = McpJson.decodeFromString(message.readText()) _onMessage.invoke(message) + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError.invoke(e) throw e diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt index b228071..b726c2f 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt @@ -1,11 +1,14 @@ package io.modelcontextprotocol.kotlin.sdk.shared +import io.kotest.assertions.throwables.shouldThrow import io.kotest.matchers.collections.shouldContainExactly +import io.kotest.matchers.collections.shouldHaveSize import io.kotest.matchers.nulls.shouldNotBeNull import io.kotest.matchers.shouldBe import io.modelcontextprotocol.kotlin.sdk.types.CustomRequest import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse import io.modelcontextprotocol.kotlin.sdk.types.McpJson @@ -24,6 +27,7 @@ import kotlinx.serialization.json.encodeToJsonElement import kotlinx.serialization.json.int import kotlinx.serialization.json.jsonObject import kotlinx.serialization.json.jsonPrimitive +import kotlin.coroutines.cancellation.CancellationException import kotlin.test.BeforeTest import kotlin.test.Test @@ -126,6 +130,36 @@ class ProtocolTest { inFlight.await() } + @Test + fun `should propagate CancellationException from notification handler without calling onError`() = runTest { + protocol.connect(transport) + + protocol.fallbackNotificationHandler = { + throw CancellationException("test cancellation") + } + + shouldThrow { + transport.deliver(JSONRPCNotification(method = "test/notification")) + } + + protocol.errors shouldHaveSize 0 + } + + @Test + fun `should report non-cancellation exception from notification handler via onError`() = runTest { + protocol.connect(transport) + + protocol.fallbackNotificationHandler = { + throw IllegalStateException("handler failed") + } + + // Non-CE exceptions are caught and reported, not propagated + transport.deliver(JSONRPCNotification(method = "test/notification")) + + protocol.errors shouldHaveSize 1 + protocol.errors[0].message shouldBe "handler failed" + } + @Test fun `should create params object when request params are null`() = runTest { protocol.connect(transport) @@ -154,6 +188,12 @@ class ProtocolTest { } private class TestProtocol : Protocol(null) { + val errors = mutableListOf() + + override fun onError(error: Throwable) { + errors.add(error) + } + override fun assertCapabilityForMethod(method: Method) { // noop } diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SSEServerTransport.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SSEServerTransport.kt index f8c6581..fcf1a46 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SSEServerTransport.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SSEServerTransport.kt @@ -84,6 +84,8 @@ public class SseServerTransport(private val endpoint: String, private val sessio } call.receiveText() + } catch (e: CancellationException) { + throw e } catch (e: Exception) { call.respondText("Invalid message: ${e.message}", status = HttpStatusCode.BadRequest) _onError.invoke(e) @@ -92,6 +94,8 @@ public class SseServerTransport(private val endpoint: String, private val sessio try { handleMessage(body) + } catch (e: CancellationException) { + throw e } catch (e: Exception) { call.respondText("Error handling message $body: ${e.message}", status = HttpStatusCode.BadRequest) return @@ -108,6 +112,8 @@ public class SseServerTransport(private val endpoint: String, private val sessio try { val parsedMessage = McpJson.decodeFromString(message) _onMessage.invoke(parsedMessage) + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError.invoke(e) throw e diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt index 9c105f9..497005a 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt @@ -28,15 +28,18 @@ import io.modelcontextprotocol.kotlin.sdk.types.RPCError import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode.REQUEST_TIMEOUT import io.modelcontextprotocol.kotlin.sdk.types.RequestId import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS +import kotlinx.coroutines.NonCancellable import kotlinx.coroutines.awaitCancellation import kotlinx.coroutines.job import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock +import kotlinx.coroutines.withContext import kotlinx.serialization.json.JsonArray import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.decodeFromJsonElement import kotlin.concurrent.atomics.AtomicBoolean import kotlin.concurrent.atomics.ExperimentalAtomicApi +import kotlin.coroutines.cancellation.CancellationException import kotlin.time.Duration import kotlin.time.Duration.Companion.milliseconds import kotlin.uuid.ExperimentalUuidApi @@ -283,17 +286,19 @@ public class StreamableHttpServerTransport(private val configuration: Configurat } override suspend fun close() { - streamMutex.withLock { - streamsMapping.values.forEach { - try { - it.session?.close() - } catch (_: Exception) { + withContext(NonCancellable) { + streamMutex.withLock { + streamsMapping.values.forEach { + try { + it.session?.close() + } catch (_: Exception) { + } } + streamsMapping.clear() + requestToStreamMapping.clear() + requestToResponseMapping.clear() + invokeOnCloseCallback() } - streamsMapping.clear() - requestToStreamMapping.clear() - requestToResponseMapping.clear() - invokeOnCloseCallback() } } @@ -420,6 +425,8 @@ public class StreamableHttpServerTransport(private val configuration: Configurat call.coroutineContext.job.invokeOnCompletion { streamsMapping.remove(streamId) } messages.forEach { message -> _onMessage(message) } + } catch (e: CancellationException) { + throw e } catch (e: Exception) { call.reject( HttpStatusCode.BadRequest, @@ -498,12 +505,14 @@ public class StreamableHttpServerTransport(private val configuration: Configurat val streamId = requestToStreamMapping[requestId] ?: return val sessionContext = streamsMapping[streamId] ?: return - try { - sessionContext.session?.close() - } catch (e: Exception) { - _onError(e) - } finally { - streamsMapping.remove(streamId) + withContext(NonCancellable) { + try { + sessionContext.session?.close() + } catch (e: Exception) { + _onError(e) + } finally { + streamsMapping.remove(streamId) + } } } @@ -554,6 +563,8 @@ public class StreamableHttpServerTransport(private val configuration: Configurat id = eventId, data = McpJson.encodeToString(message), ) + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError(IllegalStateException("Failed to replay event: ${e.message}", e)) } @@ -565,6 +576,8 @@ public class StreamableHttpServerTransport(private val configuration: Configurat streamsMapping.remove(streamId) throwable?.let { _onError(it) } } + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError(e) } @@ -658,6 +671,8 @@ public class StreamableHttpServerTransport(private val configuration: Configurat private suspend fun flushSse(session: ServerSSESession?) { try { session?.send(data = "") + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError(e) } @@ -709,6 +724,9 @@ public class StreamableHttpServerTransport(private val configuration: Configurat val eventId = configuration.eventStore?.storeEvent(streamId, message) try { session?.send(event = "message", id = eventId, data = McpJson.encodeToString(message)) + } catch (e: CancellationException) { + streamsMapping.remove(streamId) + throw e } catch (_: Exception) { streamsMapping.remove(streamId) } @@ -733,6 +751,8 @@ public class StreamableHttpServerTransport(private val configuration: Configurat retry = configuration.retryInterval?.inWholeMilliseconds, data = "", ) + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError(e) }