feat: support suspending close callbacks

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

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
         }
 

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
     }
 

kotlin-sdk-core/api/kotlin-sdk-core.api

@@ -33,8 +33,8 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/AbstractTranspor
 	public fun <init> ()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;

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() }
+            }
         }
     }
 }

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" }

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.

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<Throwable?>()
+        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()
     }
 }

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) {}
         }

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<JSONRPCMessage>(Channel.UNLIMITED)
     private var onMessageCallback: (suspend (JSONRPCMessage) -> Unit)? = null
-    private var onCloseCallback: (() -> Unit)? = null
+    private var onCloseCallback: (suspend () -> Unit)? = null
 
     val sentWithOptions = mutableListOf<Pair<JSONRPCMessage, TransportSendOptions?>>()
 
-    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
     }

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;

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<Throwable?>()
+        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()
     }
 

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()