fix!(server): #575 Rethrow CancellationException in StdioServerTransport, fix ReadBuffer

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

@@ -20,45 +20,70 @@ public class ReadBuffer {
         buffer.write(chunk)
     }
 
+    /**
+     * Reads and deserializes a JSON-RPC message from the input buffer.
+     *
+     * The method attempts to read lines from the buffer until a valid `JSONRPCMessage` is successfully
+     * deserialized. Blank lines are ignored, and if a deserialization error occurs, the method attempts
+     * to recover and process the message.
+     *
+     * Recovery involves attempting to parse from the first detected JSON object in the line when an
+     * error is encountered during deserialization.
+     *
+     * @return A deserialized `JSONRPCMessage` if successfully processed; otherwise, `null` if the
+     *         input buffer is exhausted or no valid message is found.
+     */
     public fun readMessage(): JSONRPCMessage? {
-        if (buffer.exhausted()) return null
-        var lfIndex = buffer.indexOf('\n'.code.toByte())
-        val line = when (lfIndex) {
-            -1L -> return null
+        while (true) {
+            val line = readNextLine() ?: return null
+            if (line.isBlank()) continue
 
-            0L -> {
-                buffer.skip(1)
-                return null
+            @Suppress("TooGenericExceptionCaught")
+            val message = try {
+                deserializeMessage(line)
+            } catch (e: Exception) {
+                logger.error(e) { "Failed to deserialize message from line: $line\nAttempting to recover..." }
+                tryRecover(line)
             }
-
-            else -> {
-                var skipBytes = 1
-                if (buffer[lfIndex - 1] == '\r'.code.toByte()) {
-                    lfIndex -= 1
-                    skipBytes += 1
-                }
-                val string = buffer.readString(lfIndex)
-                buffer.skip(skipBytes.toLong())
-                string
+            if (message != null) {
+                return message
             }
         }
-        try {
-            return deserializeMessage(line)
-        } catch (e: Exception) {
-            logger.error(e) { "Failed to deserialize message from line: $line\nAttempting to recover..." }
-            // if there is a non-JSON object prefix, try to parse from the first '{' onward.
-            val braceIndex = line.indexOf('{')
-            if (braceIndex != -1) {
-                val trimmed = line.substring(braceIndex)
-                try {
-                    return deserializeMessage(trimmed)
-                } catch (ignored: Exception) {
-                    logger.error(ignored) { "Deserialization failed for line: $line\nSkipping..." }
-                }
+    }
+
+    private fun readNextLine(): String? {
+        val lfIndex = if (buffer.exhausted()) -1L else buffer.indexOf('\n'.code.toByte())
+        if (lfIndex == -1L) return null
+
+        return if (lfIndex == 0L) {
+            buffer.skip(1)
+            ""
+        } else {
+            var skipBytes = 1
+            var messageLength = lfIndex
+            if (buffer[lfIndex - 1] == '\r'.code.toByte()) {
+                messageLength -= 1
+                skipBytes += 1
             }
+            val string = buffer.readString(messageLength)
+            buffer.skip(skipBytes.toLong())
+            string
         }
+    }
+
+    private fun tryRecover(line: String): JSONRPCMessage? {
+        // if there is a non-JSON object prefix, try to parse from the first '{' onward.
+        val braceIndex = line.indexOf('{')
+        if (braceIndex == -1) return null
 
-        return null
+        val trimmed = line.substring(braceIndex)
+        @Suppress("TooGenericExceptionCaught")
+        return try {
+            deserializeMessage(trimmed)
+        } catch (ignored: Exception) {
+            logger.error(ignored) { "Deserialization failed for line: $line\nSkipping..." }
+            null
+        }
     }
 
     public fun clear() {

kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt

@@ -31,6 +31,8 @@ class ReadBufferTest {
         val messageBytes = json.encodeToString(testMessage).encodeToByteArray()
         readBuffer.append(messageBytes)
         assertNull(readBuffer.readMessage())
+        readBuffer.append("\r".encodeToByteArray())
+        assertNull(readBuffer.readMessage())
 
         // Append a newline and verify message is now available
         readBuffer.append("\n".encodeToByteArray())
@@ -45,6 +47,20 @@ class ReadBufferTest {
         assertNull(readBuffer.readMessage())
     }
 
+    @Test
+    fun `skip blank line`() {
+        val readBuffer = ReadBuffer()
+        readBuffer.append(" \n".toByteArray())
+        assertNull(readBuffer.readMessage())
+    }
+
+    @Test
+    fun `skip invalid json line`() {
+        val readBuffer = ReadBuffer()
+        readBuffer.append(" {ah=oh\n".toByteArray())
+        assertNull(readBuffer.readMessage())
+    }
+
     @Test
     fun `should be reusable after clearing`() {
         val readBuffer = ReadBuffer()

kotlin-sdk-server/build.gradle.kts

@@ -23,11 +23,13 @@ kotlin {
                 implementation(libs.kotest.assertions.core)
                 implementation(libs.kotest.assertions.json)
                 implementation(libs.kotlinx.coroutines.test)
+                implementation(project(":test-utils"))
             }
         }
 
         jvmTest {
             dependencies {
+                implementation(libs.junit.jupiter.params)
                 implementation(libs.kotest.assertions.ktor)
                 implementation(libs.ktor.client.content.negotiation)
                 implementation(libs.ktor.client.logging)

kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt

@@ -7,6 +7,7 @@ import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
 import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions
 import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
 import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.NonCancellable
@@ -26,37 +27,55 @@ import kotlin.concurrent.atomics.AtomicBoolean
 import kotlin.concurrent.atomics.ExperimentalAtomicApi
 import kotlin.coroutines.CoroutineContext
 
+private const val READ_BUFFER_SIZE = 8192L
+
 /**
  * A server transport that communicates with a client via standard I/O.
  *
- * Reads from System.in and writes to System.out.
+ * Reads from input [Source] and writes to output [Sink].
+ *
+ * @constructor Creates a new instance of [StdioServerTransport].
+ * @param inputStream The input [Source] used to receive data.
+ * @param outputStream The output [Sink] used to send data.
  */
 @OptIn(ExperimentalAtomicApi::class)
 public class StdioServerTransport(private val inputStream: Source, outputStream: Sink) : AbstractTransport() {
-    private val logger = KotlinLogging.logger {}
 
+    private val logger = KotlinLogging.logger {}
     private val readBuffer = ReadBuffer()
     private val initialized: AtomicBoolean = AtomicBoolean(false)
     private var readingJob: Job? = null
     private var sendingJob: Job? = null
+    private var processingJob: Job? = null
 
     private val coroutineContext: CoroutineContext = IODispatcher + SupervisorJob()
     private val scope = CoroutineScope(coroutineContext)
     private val readChannel = Channel<ByteArray>(Channel.UNLIMITED)
     private val writeChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
-    private val outputWriter = outputStream.buffered()
+    private val outputSink = outputStream.buffered()
 
     override suspend fun start() {
         if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
             error("StdioServerTransport already started!")
         }
 
         // Launch a coroutine to read from stdin
-        readingJob = scope.launch {
+        readingJob = launchReadingJob()
+
+        // Launch a coroutine to process messages from readChannel
+        processingJob = launchProcessingJob()
+
+        // Launch a coroutine to handle message sending
+        sendingJob = launchSendingJob()
+    }
+
+    private fun launchReadingJob(): Job {
+        val job = scope.launch {
             val buf = Buffer()
+            @Suppress("TooGenericExceptionCaught")
             try {
                 while (isActive) {
-                    val bytesRead = inputStream.readAtMostTo(buf, 8192)
+                    val bytesRead = inputStream.readAtMostTo(buf, READ_BUFFER_SIZE)
                     if (bytesRead == -1L) {
                         // EOF reached
                         break
@@ -66,47 +85,73 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
                         readChannel.send(chunk)
                     }
                 }
-            } catch (e: Throwable) {
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
                 logger.error(e) { "Error reading from stdin" }
                 _onError.invoke(e)
             } finally {
                 // Reached EOF or error, close connection
                 close()
             }
         }
+        job.invokeOnCompletion { cause ->
+            logJobCompletion("Message reading", cause)
+        }
+        return job
+    }
 
-        // Launch a coroutine to process messages from readChannel
-        scope.launch {
+    private fun launchProcessingJob(): Job {
+        val job = scope.launch {
+            @Suppress("TooGenericExceptionCaught")
             try {
                 for (chunk in readChannel) {
                     readBuffer.append(chunk)
                     processReadBuffer()
                 }
-            } catch (e: Throwable) {
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
                 _onError.invoke(e)
             }
         }
+        job.invokeOnCompletion { cause ->
+            logJobCompletion("Processing", cause)
+        }
+        return job
+    }
 
-        // Launch a coroutine to handle message sending
-        sendingJob = scope.launch {
+    private fun launchSendingJob(): Job {
+        val job = scope.launch {
+            @Suppress("TooGenericExceptionCaught")
             try {
                 for (message in writeChannel) {
                     val json = serializeMessage(message)
-                    outputWriter.writeString(json)
-                    outputWriter.flush()
+                    outputSink.writeString(json)
+                    outputSink.flush()
                 }
-            } catch (e: Throwable) {
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
                 logger.error(e) { "Error writing to stdout" }
                 _onError.invoke(e)
             }
         }
+        job.invokeOnCompletion { cause ->
+            logJobCompletion("Message sending", cause)
+            if (cause is CancellationException) {
+                readingJob?.cancel(cause)
+            }
+        }
+        return job
     }
 
     private suspend fun processReadBuffer() {
+        @Suppress("TooGenericExceptionCaught")
         while (true) {
             val message = try {
                 readBuffer.readMessage()
-            } catch (e: Throwable) {
+            } catch (e: Exception) {
                 _onError.invoke(e)
                 null
             }
@@ -115,12 +160,29 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
             // Async invocation broke delivery order
             try {
                 _onMessage.invoke(message)
-            } catch (e: Throwable) {
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
                 _onError.invoke(e)
             }
         }
     }
 
+    private fun logJobCompletion(jobName: String, cause: Throwable?) {
+        when (cause) {
+            is CancellationException -> {
+            }
+
+            null -> {
+                logger.debug { "$jobName job completed" }
+            }
+
+            else -> {
+                logger.debug(cause) { "$jobName job completed exceptionally" }
+            }
+        }
+    }
+
     override suspend fun close() {
         if (!initialized.compareAndSet(expectedValue = true, newValue = false)) return
 
@@ -133,13 +195,15 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
             }.onFailure { logger.warn(it) { "Failed to close stdin" } }
 
             readingJob?.cancel()
-
             readChannel.close()
+
+            processingJob?.cancelAndJoin()
+
             readBuffer.clear()
 
             runCatching {
-                outputWriter.flush()
-                outputWriter.close()
+                outputSink.flush()
+                outputSink.close()
             }.onFailure { logger.warn(it) { "Failed to close stdout" } }
 
             invokeOnCloseCallback()