diff --git a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt index 3352104fb..73b6bb2a9 100644 --- a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt +++ b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt @@ -225,7 +225,9 @@ public open class Client(private val clientInfo: Implementation, options: Client notification(InitializedNotification()) enableConcurrentDispatch() } catch (error: Throwable) { - logger.error(error) { "Failed to initialize client: ${error.message}" } + if (error !is CancellationException) { + logger.error(error) { "Failed to initialize client: ${error.message}" } + } close() when (error) { 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 0295eb9a3..04f399187 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 @@ -156,6 +156,8 @@ public class SseClientTransport( } endpoint.complete(endpointUrl) logger.debug { "Client connected to endpoint: $endpointUrl" } + } catch (e: CancellationException) { + throw e } catch (e: Throwable) { _onError(e) endpoint.completeExceptionally(e) @@ -179,6 +181,8 @@ public class SseClientTransport( if (::session.isInitialized) session.cancel() if (::scope.isInitialized) scope.cancel() endpoint.cancel() + } catch (e: CancellationException) { + throw e } 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 5e990eea0..4daab4651 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 @@ -179,11 +179,9 @@ public class StdioClientTransport @JvmOverloads public constructor( val errorSeverity = classifyStderr(event.message) when (errorSeverity) { FATAL -> { - runCatching { - _onError( - McpException(INTERNAL_ERROR, "Message in StdErr: ${event.message}"), - ) - } + invokeOnErrorCallback( + McpException(INTERNAL_ERROR, "Message in StdErr: ${event.message}"), + ) stopProcessing("Fatal STDERR message received") } @@ -212,7 +210,7 @@ public class StdioClientTransport @JvmOverloads public constructor( } is Event.IOErrorEvent -> { - runCatching { _onError(event.cause) } + invokeOnErrorCallback(event.cause) stopProcessing("IO Error", event.cause) } } @@ -265,11 +263,11 @@ public class StdioClientTransport @JvmOverloads public constructor( sink.flush() } catch (e: SerializationException) { logger.warn(e) { "Can't serialize message" } - runCatching { _onError(McpException(INTERNAL_ERROR, "Serialization error")) } + invokeOnErrorCallback(McpException(INTERNAL_ERROR, "Serialization error")) mainScope.stopProcessing("Can't serialize message", e) } catch (e: IOException) { logger.warn(e) { "Can't send message" } - runCatching { _onError(McpException(CONNECTION_CLOSED, "Can't send message. Connection closed")) } + invokeOnErrorCallback(McpException(CONNECTION_CLOSED, "Can't send message. Connection closed")) mainScope.stopProcessing("Write I/O failed", e) } } @@ -281,7 +279,7 @@ public class StdioClientTransport @JvmOverloads public constructor( throw e } catch (e: Throwable) { logger.error(e) { "Error processing message." } - runCatching { _onError.invoke(e) } + invokeOnErrorCallback(e) } } diff --git a/kotlin-sdk-core/api/kotlin-sdk-core.api b/kotlin-sdk-core/api/kotlin-sdk-core.api index b3e88d84d..11e6d4e58 100644 --- a/kotlin-sdk-core/api/kotlin-sdk-core.api +++ b/kotlin-sdk-core/api/kotlin-sdk-core.api @@ -34,6 +34,7 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/AbstractTranspor protected final fun get_onError ()Lkotlin/jvm/functions/Function1; protected final fun get_onMessage ()Lkotlin/jvm/functions/Function2; protected final fun invokeOnCloseCallback ()V + protected final fun invokeOnErrorCallback (Ljava/lang/Throwable;)V public fun onClose (Lkotlin/jvm/functions/Function0;)V public fun onError (Lkotlin/jvm/functions/Function1;)V public fun onMessage (Lkotlin/jvm/functions/Function2;)V diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport.kt index a73828979..b46d09a49 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport.kt @@ -1,6 +1,7 @@ package io.modelcontextprotocol.kotlin.sdk.shared import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.utils.runCatchingCancellable import kotlinx.coroutines.CompletableDeferred import kotlin.concurrent.atomics.AtomicBoolean import kotlin.concurrent.atomics.ExperimentalAtomicApi @@ -67,7 +68,15 @@ public abstract class AbstractTransport : Transport { */ protected fun invokeOnCloseCallback() { if (onCloseCalled.compareAndSet(expectedValue = false, newValue = true)) { - runCatching { _onClose() } + runCatchingCancellable { _onClose() } } } + + /** + * Reports [error] through the `_onError` callback, swallowing any [Throwable] the callback + * raises. A [kotlin.coroutines.cancellation.CancellationException] propagates instead. + */ + protected fun invokeOnErrorCallback(error: Throwable) { + runCatchingCancellable { _onError(error) } + } } 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 908e524ce..e04714926 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 @@ -870,6 +870,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio withContext(NonCancellable) { try { cancelPending(timeoutError, notifyPeerOnCancel) + } catch (_: CancellationException) { } catch (e: Throwable) { logger.warn(e) { "Failed to notify peer about timed-out request" } onError(e) @@ -887,6 +888,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio withContext(NonCancellable) { try { cancelPending(cause, notifyPeerOnCancel) + } catch (_: CancellationException) { } catch (e: Throwable) { logger.warn(e) { "Failed to notify peer about cancelled request" } onError(e) 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 baf37233e..820864800 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 @@ -86,7 +86,7 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { @OptIn(InternalCoroutinesApi::class) session.coroutineContext.job.invokeOnCompletion { - if (it != null) { + if (it != null && it !is CancellationException) { _onError.invoke(it) } else { invokeOnCloseCallback() diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/utils/RunCatchingCancellable.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/utils/RunCatchingCancellable.kt new file mode 100644 index 000000000..3484453b4 --- /dev/null +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/utils/RunCatchingCancellable.kt @@ -0,0 +1,15 @@ +package io.modelcontextprotocol.kotlin.sdk.utils + +import kotlin.coroutines.cancellation.CancellationException + +/** + * Like [runCatching], but re-throws [CancellationException] instead of capturing it into a failed + * [Result]: swallowing cancellation would let a cancelled coroutine keep running. + */ +internal inline fun runCatchingCancellable(block: () -> T): Result = try { + Result.success(block()) +} catch (e: CancellationException) { + throw e +} catch (e: Throwable) { + Result.failure(e) +} diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/utils/RunCatchingCancellableTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/utils/RunCatchingCancellableTest.kt new file mode 100644 index 000000000..bc4c3559c --- /dev/null +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/utils/RunCatchingCancellableTest.kt @@ -0,0 +1,38 @@ +package io.modelcontextprotocol.kotlin.sdk.utils + +import io.kotest.matchers.shouldBe +import io.kotest.matchers.types.shouldBeInstanceOf +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.cancelAndJoin +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlin.test.Test + +class RunCatchingCancellableTest { + + @Test + fun `should capture a regular throwable as failure`() { + val result = runCatchingCancellable { throw IllegalStateException("boom") } + + result.exceptionOrNull().shouldBeInstanceOf().message shouldBe "boom" + } + + @Test + fun `should not let a cancelled coroutine continue past the call`() = runTest { + val entered = CompletableDeferred() + var reachedAfterCall = false + + val job = launch { + runCatchingCancellable { + entered.complete(Unit) + awaitCancellation() + } + reachedAfterCall = true + } + entered.await() + job.cancelAndJoin() + + reachedAfterCall shouldBe false + } +} 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 dcba84873..e85d71bfa 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 @@ -664,6 +664,8 @@ public class StreamableHttpServerTransport(private val configuration: Configurat withContext(NonCancellable) { try { sessionContext.session?.close() + } catch (e: CancellationException) { + throw e } catch (e: Exception) { _onError(e) } finally { @@ -729,7 +731,7 @@ public class StreamableHttpServerTransport(private val configuration: Configurat session.coroutineContext.job.invokeOnCompletion { throwable -> streamsMapping.remove(streamId) - throwable?.let { _onError(it) } + if (throwable != null && throwable !is CancellationException) _onError(throwable) } } catch (e: CancellationException) { throw e