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

## Rethrow CancellationException in StdioServerTransport, fix ReadBuffer

### Changes
1. StdioServerTransport
- Replace generic `Throwable` catches with specific
`CancellationException` and `Exception` for clarity and correctness.
**breaking change** `Throwable` is not handled any more!
    - Refactor launching jobs to separate methods
- Introduce `READ_BUFFER_SIZE` constant to replace inline buffer size
literals.
- Add suppress annotations for clearer intent and constructor
documentation for better usability.
    - Extend StdioServerTransportTest

2. ReadBuffer

Previously, readMessage() returned null after consuming an unparseable
line even when more complete lines existed in the buffer. This caused
valid messages following a bad line in the same chunk to be silently
dropped until the next chunk arrived (or forever, in tests).
    
Fix the loop so null means only "no complete line available", not
"encountered a parse failure". Blank/whitespace-only lines are now
silently skipped via isBlank() rather than forwarded to the deserializer
and logged as errors.
    
    Refactor the method into three focused helpers:
     - readMessage() — outer loop over lines
- readLine() — consume the next newline-delimited line from the buffer
     - tryRecover()  — attempt deserialization from the first '{' onward

3. Add utility method `runIntegrationTest` for integration testing
(`runBlocking` + `withTimeout`). Use it in StdioServerTransportTest

4. Store processing Job reference and `calcelAndJoin()` it on close.
Partially addressing #574

## Motivation and Context

See #575, #564, #242 

## How Has This Been Tested?
Unit test added

## Breaking Changes

Semantic change: `Throwable` and CancellationException are not handled
any more!

## Types of changes

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update

## Checklist
<!-- Go over all the following points, and put an `x` in all the boxes
that apply. -->
- [x] I have read the [MCP
Documentation](https://modelcontextprotocol.io)
- [x] My code follows the repository's style guidelines
- [x] New and existing tests pass locally
- [ ] I have added appropriate error handling
- [ ] I have added or updated documentation as needed

## Additional context
<!-- Add any other context, implementation notes, or design decisions
-->

Author: kpavlov

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,6 +85,8 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
                         readChannel.send(chunk)
                     }
                 }
+            } catch (e: CancellationException) {
+                throw e
             } catch (e: Throwable) {
                 logger.error(e) { "Error reading from stdin" }
                 _onError.invoke(e)
@@ -74,35 +95,59 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
                 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: CancellationException) {
+                throw e
             } catch (e: Throwable) {
                 _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: CancellationException) {
+                throw e
             } catch (e: Throwable) {
                 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()
@@ -115,12 +160,30 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
             // Async invocation broke delivery order
             try {
                 _onMessage.invoke(message)
+            } catch (e: CancellationException) {
+                throw e
             } catch (e: Throwable) {
+                logger.error(e) { "Error processing message" }
                 _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 +196,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()