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..9d065b9f9 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 @@ -258,7 +258,7 @@ public class StdioClientTransport @JvmOverloads public constructor( } } - private fun sendOutboundMessage(message: JSONRPCMessage, sink: Sink, mainScope: CoroutineScope) { + private suspend fun sendOutboundMessage(message: JSONRPCMessage, sink: Sink, mainScope: CoroutineScope) { try { val json = serializeMessage(message) sink.writeString(json) @@ -285,7 +285,7 @@ public class StdioClientTransport @JvmOverloads public constructor( } } - private fun CoroutineScope.stopProcessing(reason: String, cause: Throwable? = null) { + private suspend fun CoroutineScope.stopProcessing(reason: String, cause: Throwable? = null) { sendChannel.close() // Stop accepting new messages invokeOnCloseCallback() cancel(reason, cause) // cancel current coroutine context diff --git a/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientAssertCapabilityTest.kt b/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientAssertCapabilityTest.kt index e9379a7f0..30e82df03 100644 --- a/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientAssertCapabilityTest.kt +++ b/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientAssertCapabilityTest.kt @@ -135,7 +135,7 @@ class ClientAssertCapabilityTest { */ private class CapabilitiesTransport(private val serverCapabilities: ServerCapabilities) : Transport { private var onMessageBlock: (suspend (JSONRPCMessage) -> Unit)? = null - private var onCloseBlock: (() -> Unit)? = null + private var onCloseBlock: (suspend () -> Unit)? = null override suspend fun start() = Unit @@ -162,7 +162,7 @@ class ClientAssertCapabilityTest { onMessageBlock = block } - override fun onClose(block: () -> Unit) { + override fun onClose(block: suspend () -> Unit) { onCloseBlock = block } diff --git a/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/MockTransport.kt b/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/MockTransport.kt index d7a48efdf..50f887a88 100644 --- a/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/MockTransport.kt +++ b/kotlin-sdk-client/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/MockTransport.kt @@ -21,7 +21,7 @@ class MockTransport : Transport { suspend fun getReceivedMessages() = mutex.withLock { _receivedMessages.toList() } private var onMessageBlock: (suspend (JSONRPCMessage) -> Unit)? = null - private var onCloseBlock: (() -> Unit)? = null + private var onCloseBlock: (suspend () -> Unit)? = null private var onErrorBlock: ((Throwable) -> Unit)? = null override suspend fun start() = Unit @@ -81,7 +81,7 @@ class MockTransport : Transport { } } - override fun onClose(block: () -> Unit) { + override fun onClose(block: suspend () -> Unit) { onCloseBlock = block } diff --git a/kotlin-sdk-core/api/kotlin-sdk-core.api b/kotlin-sdk-core/api/kotlin-sdk-core.api index b3e88d84d..ec07442c3 100644 --- a/kotlin-sdk-core/api/kotlin-sdk-core.api +++ b/kotlin-sdk-core/api/kotlin-sdk-core.api @@ -33,8 +33,8 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/AbstractTranspor public fun ()V protected final fun get_onError ()Lkotlin/jvm/functions/Function1; protected final fun get_onMessage ()Lkotlin/jvm/functions/Function2; - protected final fun invokeOnCloseCallback ()V - public fun onClose (Lkotlin/jvm/functions/Function0;)V + protected final fun invokeOnCloseCallback (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public fun onClose (Lkotlin/jvm/functions/Function1;)V public fun onError (Lkotlin/jvm/functions/Function1;)V public fun onMessage (Lkotlin/jvm/functions/Function2;)V } @@ -70,7 +70,7 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol { public final fun getTransport ()Lio/modelcontextprotocol/kotlin/sdk/shared/Transport; public final fun notification (Lio/modelcontextprotocol/kotlin/sdk/types/Notification;Lio/modelcontextprotocol/kotlin/sdk/types/RequestId;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public static synthetic fun notification$default (Lio/modelcontextprotocol/kotlin/sdk/shared/Protocol;Lio/modelcontextprotocol/kotlin/sdk/types/Notification;Lio/modelcontextprotocol/kotlin/sdk/types/RequestId;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; - public fun onClose ()V + public fun onClose (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public fun onError (Ljava/lang/Throwable;)V protected fun onInitializedNotification ()V public final fun removeNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V @@ -146,7 +146,7 @@ public final class io/modelcontextprotocol/kotlin/sdk/shared/TooLongFrameExcepti public abstract interface class io/modelcontextprotocol/kotlin/sdk/shared/Transport { public abstract fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; - public abstract fun onClose (Lkotlin/jvm/functions/Function0;)V + public abstract fun onClose (Lkotlin/jvm/functions/Function1;)V public abstract fun onError (Lkotlin/jvm/functions/Function1;)V public abstract fun onMessage (Lkotlin/jvm/functions/Function2;)V public abstract fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; 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..c2b6f43d6 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 @@ -2,6 +2,8 @@ package io.modelcontextprotocol.kotlin.sdk.shared import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.withContext import kotlin.concurrent.atomics.AtomicBoolean import kotlin.concurrent.atomics.ExperimentalAtomicApi @@ -15,7 +17,7 @@ import kotlin.concurrent.atomics.ExperimentalAtomicApi @OptIn(ExperimentalAtomicApi::class) public abstract class AbstractTransport : Transport { private val onCloseCalled = AtomicBoolean(false) - private var _onClose: (() -> Unit) = {} + private var _onClose: suspend () -> Unit = {} protected var _onError: ((Throwable) -> Unit) = {} private set @@ -27,8 +29,8 @@ public abstract class AbstractTransport : Transport { } private set - override fun onClose(block: () -> Unit) { - val old = _onClose + override fun onClose(block: suspend () -> Unit) { + val old: suspend () -> Unit = _onClose _onClose = { old() block() @@ -60,14 +62,17 @@ public abstract class AbstractTransport : Transport { /** * Invokes the `_onClose` callback if it has not been already triggered. * - * This method ensures the `_onClose` callback is executed only once by utilizing + * This suspending method ensures the `_onClose` callback is awaited and executed only once by utilizing * an atomic flag (`onCloseCalled`). If the callback has already been executed, - * the method does nothing. Any exceptions thrown during the execution of the - * `_onClose` callback are caught and suppressed. + * the method does nothing. The callback runs in [NonCancellable] so suspending cleanup can + * finish after its owning transport scope has stopped. Any exceptions thrown during the + * execution of the `_onClose` callback are caught and suppressed. */ - protected fun invokeOnCloseCallback() { + protected suspend fun invokeOnCloseCallback() { if (onCloseCalled.compareAndSet(expectedValue = false, newValue = true)) { - runCatching { _onClose() } + withContext(NonCancellable) { + runCatching { _onClose() } + } } } } 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..fbc6f0b7a 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 @@ -297,9 +297,10 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio /** * Callback for when the connection is closed for any reason. * - * This is invoked when close() is called as well. + * This suspending hook is invoked and awaited when [close] is called as well. Subclasses may + * override it to finish asynchronous cleanup before the close path completes. */ - public open fun onClose() {} + public open suspend fun onClose() {} /** * Callback for when an error occurs. @@ -429,7 +430,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } } - private fun doClose(connection: Connection) { + private suspend fun doClose(connection: Connection) { // A stale onClose from a previous transport must not tear down the successor connection. if (!connectionRef.compareAndSet(connection, null)) { logger.trace { "Ignoring close signal from a stale transport" } diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt index a89c10291..7704b1450 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt @@ -35,9 +35,13 @@ public interface Transport { /** * Callback for when the connection is closed for any reason. * - * This should be invoked when close() is called as well. + * This should be invoked and awaited when [close] is called as well. The callback may suspend + * to finish asynchronous resource cleanup before the transport completes its close path. + * Multiple callbacks are invoked in registration order. + * + * @param block suspending cleanup to invoke when the connection closes */ - public fun onClose(block: () -> Unit) + public fun onClose(block: suspend () -> Unit) /** * Callback for when an error occurs. 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..1cbcf83ef 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 @@ -7,10 +7,13 @@ import io.ktor.websocket.close import io.ktor.websocket.readText import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.CoroutineName import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.InternalCoroutinesApi +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.Job import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.cancel import kotlinx.coroutines.channels.ClosedReceiveChannelException import kotlinx.coroutines.job import kotlinx.coroutines.launch @@ -34,6 +37,7 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { } private val initialized: AtomicBoolean = AtomicBoolean(false) + private var closeWatcherJob: Job? = null /** * The WebSocket session used for communication. @@ -84,12 +88,21 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { } } - @OptIn(InternalCoroutinesApi::class) - session.coroutineContext.job.invokeOnCompletion { - if (it != null) { - _onError.invoke(it) - } else { - invokeOnCloseCallback() + val sessionCompletion = CompletableDeferred() + session.coroutineContext.job.invokeOnCompletion { sessionCompletion.complete(it) } + closeWatcherJob = scope.launch( + context = CoroutineName("WebSocketMcpTransport.close#${hashCode()}"), + start = CoroutineStart.UNDISPATCHED, + ) { + try { + val cause = sessionCompletion.await() + if (cause != null) { + _onError.invoke(cause) + } else { + invokeOnCloseCallback() + } + } finally { + scope.cancel() } } } @@ -111,5 +124,6 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { logger.debug { "Closing websocket session" } session.close() session.coroutineContext.job.join() + closeWatcherJob?.join() } } 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 2846cccaa..1423e6840 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 @@ -315,7 +315,7 @@ class ProtocolTest { override suspend fun start(): Unit = error("boom") override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {} override suspend fun close() {} - override fun onClose(block: () -> Unit) {} + override fun onClose(block: suspend () -> Unit) {} override fun onError(block: (Throwable) -> Unit) {} override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {} } @@ -333,7 +333,7 @@ class ProtocolTest { override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?): Unit = throw IllegalStateException("send failed") override suspend fun close() {} - override fun onClose(block: () -> Unit) {} + override fun onClose(block: suspend () -> Unit) {} override fun onError(block: (Throwable) -> Unit) {} override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {} } diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt index 9e8dd6093..a91f998d6 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt @@ -51,11 +51,11 @@ internal class TestProtocol(options: ProtocolOptions? = null) : Protocol(options internal class RecordingTransport : Transport { val sentMessages = Channel(Channel.UNLIMITED) private var onMessageCallback: (suspend (JSONRPCMessage) -> Unit)? = null - private var onCloseCallback: (() -> Unit)? = null + private var onCloseCallback: (suspend () -> Unit)? = null val sentWithOptions = mutableListOf>() - var closeCallback: (() -> Unit)? = null + var closeCallback: (suspend () -> Unit)? = null private set override suspend fun start() { @@ -71,7 +71,7 @@ internal class RecordingTransport : Transport { onCloseCallback?.invoke() } - override fun onClose(block: () -> Unit) { + override fun onClose(block: suspend () -> Unit) { closeCallback = block onCloseCallback = block } diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index d8e4c7a5e..e4ecb8ec9 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -145,7 +145,7 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server { public final fun getTools ()Ljava/util/Map; public final fun listRoots (Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public static synthetic fun listRoots$default (Lio/modelcontextprotocol/kotlin/sdk/server/Server;Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; - public final fun onClose (Lkotlin/jvm/functions/Function0;)V + public final fun onClose (Lkotlin/jvm/functions/Function1;)V public final fun onConnect (Lkotlin/jvm/functions/Function0;)V public final fun ping (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun removeNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V @@ -197,8 +197,8 @@ public class io/modelcontextprotocol/kotlin/sdk/server/ServerSession : io/modelc public fun hashCode ()I public final fun listRoots (Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public static synthetic fun listRoots$default (Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; - public fun onClose ()V - public final fun onClose (Lkotlin/jvm/functions/Function0;)V + public fun onClose (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public final fun onClose (Lkotlin/jvm/functions/Function1;)V public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V protected fun onInitializedNotification ()V public final fun ping (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; 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 f525ec29e..6e29aa06d 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 @@ -11,8 +11,16 @@ import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.McpJson -import kotlinx.coroutines.InternalCoroutinesApi +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineName +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.Job +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.cancel +import kotlinx.coroutines.cancelAndJoin import kotlinx.coroutines.job +import kotlinx.coroutines.launch import kotlin.concurrent.atomics.AtomicBoolean import kotlin.concurrent.atomics.ExperimentalAtomicApi import kotlin.coroutines.cancellation.CancellationException @@ -42,6 +50,10 @@ public class SseServerTransport( } private val initialized: AtomicBoolean = AtomicBoolean(false) + private val closeWatcherScope: CoroutineScope by lazy { + CoroutineScope(session.coroutineContext + SupervisorJob()) + } + private var closeWatcherJob: Job? = null /** Unique identifier for this transport session, generated randomly on creation. */ @OptIn(ExperimentalUuidApi::class) @@ -65,12 +77,21 @@ public class SseServerTransport( data = "${endpoint.encodeURLPath()}?$SESSION_ID_PARAM=$sessionId", ) - @OptIn(InternalCoroutinesApi::class) - session.coroutineContext.job.invokeOnCompletion { - if (it != null && it !is CancellationException) { - _onError.invoke(it) - } else { - invokeOnCloseCallback() + val sessionCompletion = CompletableDeferred() + session.coroutineContext.job.invokeOnCompletion { sessionCompletion.complete(it) } + closeWatcherJob = closeWatcherScope.launch( + context = CoroutineName("SseServerTransport.close#$sessionId"), + start = CoroutineStart.UNDISPATCHED, + ) { + try { + val cause = sessionCompletion.await() + if (cause != null && cause !is CancellationException) { + _onError.invoke(cause) + } else { + invokeOnCloseCallback() + } + } finally { + closeWatcherScope.cancel() } } } @@ -137,6 +158,7 @@ public class SseServerTransport( override suspend fun close() { session.close() + closeWatcherJob?.cancelAndJoin() invokeOnCloseCallback() } diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt index 51a238918..21d545a07 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt @@ -130,7 +130,7 @@ public open class Server( private var _onConnect: (() -> Unit) = {} - private var _onClose: () -> Unit = {} + private var _onClose: suspend () -> Unit = {} @OptIn(ExperimentalTime::class) private val notificationService = FeatureNotificationService() @@ -195,7 +195,10 @@ public open class Server( block(this) } - /** Closes this server, shutting down the notification service and all active sessions. */ + /** + * Closes this server, shutting down the notification service and all active sessions before + * awaiting callbacks registered with [onClose]. + */ public suspend fun close() { logger.debug { "Closing MCP server" } notificationService.close() @@ -290,10 +293,15 @@ public open class Server( } /** - * Registers a callback to be invoked when the server connection is closing. + * Registers a callback to be invoked and awaited when the server is closing. + * + * Multiple callbacks are invoked in registration order. A callback may suspend while it + * finishes asynchronous cleanup. + * + * @param block suspending cleanup to invoke when the server closes */ - public fun onClose(block: () -> Unit) { - val old = _onClose + public fun onClose(block: suspend () -> Unit) { + val old: suspend () -> Unit = _onClose _onClose = { old() block() diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt index f20d3e122..8d40f5d67 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt @@ -66,7 +66,7 @@ public open class ServerSession( private var _onInitialized: (() -> Unit) = {} - private var _onClose: () -> Unit = {} + private var _onClose: suspend () -> Unit = {} private val _clientCapabilities: AtomicRef = atomic(null) private val _clientVersion: AtomicRef = atomic(null) @@ -130,10 +130,15 @@ public open class ServerSession( } /** - * Registers a callback to be invoked when the server session is closing. + * Registers a callback to be invoked and awaited when the server session is closing. + * + * Multiple callbacks are invoked in registration order. A callback may suspend while it + * finishes asynchronous cleanup. + * + * @param block suspending cleanup to invoke when the session closes */ - public fun onClose(block: () -> Unit) { - val old = _onClose + public fun onClose(block: suspend () -> Unit) { + val old: suspend () -> Unit = _onClose _onClose = { old() block() @@ -141,9 +146,9 @@ public open class ServerSession( } /** - * Called when the server session is closing. + * Called and awaited when the server session is closing. */ - override fun onClose() { + override suspend fun onClose() { logger.debug { "Server connection closing" } _onClose() } diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt index 9261be257..aabced209 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt @@ -331,7 +331,7 @@ public class StdioServerTransport private constructor( } } - private fun transitionToStoppedNaturally() { + private suspend fun transitionToStoppedNaturally() { if (!state.compareAndSet(State.Operational, State.Stopped)) return runCatching { input.close() } .onFailure { logger.warn(it) { "Failed to close input source" } } diff --git a/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerCloseTest.kt b/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerCloseTest.kt new file mode 100644 index 000000000..92b85e423 --- /dev/null +++ b/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerCloseTest.kt @@ -0,0 +1,96 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.kotest.matchers.collections.shouldContainExactly +import io.kotest.matchers.shouldBe +import io.modelcontextprotocol.kotlin.sdk.shared.Transport +import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlin.test.Test + +class ServerCloseTest { + @Test + fun `session close awaits suspending cleanup`() = runTest { + val session = ServerSession( + serverInfo = Implementation(name = "test-server", version = "1.0"), + options = testServerOptions(), + instructions = null, + ) + val transport = CloseCallbackTransport() + session.connect(transport) + val releaseCleanup = CompletableDeferred() + val cleanupStarted = CompletableDeferred() + val events = mutableListOf() + + session.onClose { + events += "started" + cleanupStarted.complete(Unit) + releaseCleanup.await() + events += "finished" + } + + val closeJob = launch { session.close() } + + cleanupStarted.await() + closeJob.isActive shouldBe true + events shouldContainExactly listOf("started") + + releaseCleanup.complete(Unit) + closeJob.join() + events shouldContainExactly listOf("started", "finished") + } + + @Test + fun `server close awaits suspending cleanup`() = runTest { + val server = Server( + serverInfo = Implementation(name = "test-server", version = "1.0"), + options = testServerOptions(), + ) + val releaseCleanup = CompletableDeferred() + val cleanupStarted = CompletableDeferred() + val events = mutableListOf() + + server.onClose { + events += "started" + cleanupStarted.complete(Unit) + releaseCleanup.await() + events += "finished" + } + + val closeJob = launch { server.close() } + + cleanupStarted.await() + closeJob.isActive shouldBe true + events shouldContainExactly listOf("started") + + releaseCleanup.complete(Unit) + closeJob.join() + events shouldContainExactly listOf("started", "finished") + } + + private fun testServerOptions() = ServerOptions(capabilities = ServerCapabilities()) + + private class CloseCallbackTransport : Transport { + private var closeCallback: suspend () -> Unit = {} + + override suspend fun start() = Unit + + override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) = Unit + + override suspend fun close() { + closeCallback() + } + + override fun onClose(block: suspend () -> Unit) { + closeCallback = block + } + + override fun onError(block: (Throwable) -> Unit) = Unit + + override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) = Unit + } +} diff --git a/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionAssertCapabilityTest.kt b/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionAssertCapabilityTest.kt index a172e965b..142dfe003 100644 --- a/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionAssertCapabilityTest.kt +++ b/kotlin-sdk-server/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionAssertCapabilityTest.kt @@ -115,7 +115,7 @@ class ServerSessionAssertCapabilityTest { */ private class InitializeReplayTransport(private val clientCapabilities: ClientCapabilities) : Transport { private var onMessageBlock: (suspend (JSONRPCMessage) -> Unit)? = null - private var onCloseBlock: (() -> Unit)? = null + private var onCloseBlock: (suspend () -> Unit)? = null override suspend fun start() { val initializeRequest = InitializeRequest( @@ -138,7 +138,7 @@ class ServerSessionAssertCapabilityTest { onMessageBlock = block } - override fun onClose(block: () -> Unit) { + override fun onClose(block: suspend () -> Unit) { onCloseBlock = block } diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SseServerTransportTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SseServerTransportTest.kt index a0756f342..67c95f338 100644 --- a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SseServerTransportTest.kt +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SseServerTransportTest.kt @@ -16,12 +16,41 @@ import io.ktor.server.sse.SSE import io.ktor.server.sse.ServerSSESession import io.ktor.server.testing.testApplication import io.ktor.utils.io.readLine +import io.mockk.every import io.mockk.mockk +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withTimeout import java.util.concurrent.atomic.AtomicBoolean import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds class SseServerTransportTest : AbstractKtorExtensionsTest() { + @Test + fun `close awaits suspending callback without waiting for the session job`() = runTest { + val session = mockk(relaxed = true) + every { session.coroutineContext } returns currentCoroutineContext() + val transport = SseServerTransport("/messages", session) + val callbackStarted = CompletableDeferred() + val releaseCallback = CompletableDeferred() + + transport.onClose { + callbackStarted.complete(Unit) + releaseCallback.await() + } + transport.start() + + val closeJob = launch { transport.close() } + withTimeout(1.seconds) { callbackStarted.await() } + closeJob.isActive shouldBe true + + releaseCallback.complete(Unit) + closeJob.join() + } + @Test fun `handlePostMessage on a not-started transport does not deliver the message`() = testApplication { // A registered onMessage callback opens the delivery gate, so if the not-initialized branch diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransportTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransportTest.kt index c6370a1a6..b9da0615c 100644 --- a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransportTest.kt +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransportTest.kt @@ -1,6 +1,5 @@ package io.modelcontextprotocol.kotlin.sdk.server -import io.kotest.assertions.nondeterministic.eventually import io.kotest.assertions.throwables.shouldThrow import io.kotest.assertions.withClue import io.kotest.matchers.collections.shouldContain @@ -109,6 +108,41 @@ class StdioServerTransportTest { assertTrue(didClose, "Should have closed after calling close()") } + @Test + fun `should await suspending onClose callbacks in registration order exactly once`() = runIntegrationTest { + val server = StdioServerTransport(input = bufferedInput, output = printOutput) + val firstCallbackStarted = CompletableDeferred() + val releaseFirstCallback = CompletableDeferred() + val events = mutableListOf() + + server.onError { throw it } + server.onMessage {} + server.onClose { + events += "first-started" + firstCallbackStarted.complete(Unit) + releaseFirstCallback.await() + events += "first-finished" + } + server.onClose { + delay(1.milliseconds) + events += "second" + } + + server.start() + val closeJob = launch { server.close() } + + firstCallbackStarted.await() + closeJob.isActive shouldBe true + events shouldBe listOf("first-started") + + releaseFirstCallback.complete(Unit) + closeJob.join() + events shouldBe listOf("first-started", "first-finished", "second") + + server.close() + events shouldBe listOf("first-started", "first-finished", "second") + } + @Test fun `should not read until started`() = runIntegrationTest { val server = StdioServerTransport(input = bufferedInput, output = printOutput) @@ -208,15 +242,16 @@ class StdioServerTransportTest { val server = StdioServerTransport(input = bufferedInput, output = printOutput) val didClose = CompletableDeferred() server.onError { throw it } - server.onClose { didClose.complete(Unit) } + server.onClose { + delay(1.milliseconds) + didClose.complete(Unit) + } server.onMessage {} server.start() inputWriter.close() // signal EOF to the reading loop - eventually(2.seconds) { - didClose.isCompleted shouldBe true - } + didClose.await() } @Test