From 2875ca21847d68f0460235b64dc5d48594a5c727 Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Sun, 1 Mar 2026 23:32:04 +0200 Subject: [PATCH 1/7] fix!: Handle CancellationException, add buffer size constant in 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 --- kotlin-sdk-server/build.gradle.kts | 1 + .../kotlin/sdk/server/StdioServerTransport.kt | 69 +++++- .../sdk/server/StdioServerTransportTest.kt | 204 ++++++++++++++++-- 3 files changed, 244 insertions(+), 30 deletions(-) diff --git a/kotlin-sdk-server/build.gradle.kts b/kotlin-sdk-server/build.gradle.kts index 9d0bddff8..134b36869 100644 --- a/kotlin-sdk-server/build.gradle.kts +++ b/kotlin-sdk-server/build.gradle.kts @@ -28,6 +28,7 @@ kotlin { jvmTest { dependencies { + implementation(libs.junit.jupiter.params) implementation(libs.kotest.assertions.ktor) implementation(libs.ktor.client.content.negotiation) implementation(libs.ktor.client.logging) 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 b3ec4d650..9a4eee884 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 @@ -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 @@ -30,10 +31,18 @@ import kotlin.coroutines.CoroutineContext * A server transport that communicates with a client via standard I/O. * * Reads from System.in and writes to System.out. + * + * @constructor Creates a new instance of [StdioServerTransport]. + * @param inputStream The input stream used to receive data. + * @param outputStream The output stream used to send data. */ @OptIn(ExperimentalAtomicApi::class) public class StdioServerTransport(private val inputStream: Source, outputStream: Sink) : AbstractTransport() { - private val logger = KotlinLogging.logger {} + + private companion object { + private val logger = KotlinLogging.logger {} + private const val READ_BUFFER_SIZE = 8192L + } private val readBuffer = ReadBuffer() private val initialized: AtomicBoolean = AtomicBoolean(false) @@ -46,17 +55,29 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: private val writeChannel = Channel(Channel.UNLIMITED) private val outputWriter = outputStream.buffered() + @Suppress("TooGenericExceptionCaught", "ThrowsCount") 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 + launchProcessingJob() + + // Launch a coroutine to handle message sending + sendingJob = launchSendingJob() + } + + @Suppress("TooGenericExceptionCaught") + private fun launchReadingJob(): Job { + val job = scope.launch { val buf = Buffer() try { while (isActive) { - val bytesRead = inputStream.readAtMostTo(buf, 8192) + val bytesRead = inputStream.readAtMostTo(buf, READ_BUFFER_SIZE) if (bytesRead == -1L) { // EOF reached break @@ -66,7 +87,9 @@ 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 { @@ -74,39 +97,61 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: close() } } + job.invokeOnCompletion { cause -> + logger.debug(cause) { "Message reading job completed with cause: $cause" } + } + return job + } - // Launch a coroutine to process messages from readChannel + @Suppress("TooGenericExceptionCaught") + private fun launchProcessingJob() { scope.launch { try { for (chunk in readChannel) { readBuffer.append(chunk) processReadBuffer() } - } catch (e: Throwable) { + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { _onError.invoke(e) } } + } - // Launch a coroutine to handle message sending - sendingJob = scope.launch { + @Suppress("TooGenericExceptionCaught") + private fun launchSendingJob(): Job { + val job = scope.launch { try { for (message in writeChannel) { val json = serializeMessage(message) outputWriter.writeString(json) outputWriter.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 -> + logger.debug(cause) { "Message sending job completed with cause: $cause" } + if (cause is CancellationException) { + readingJob?.cancel(cause) + } + } + return job } + @Suppress("TooGenericExceptionCaught") private suspend fun processReadBuffer() { while (true) { val message = try { readBuffer.readMessage() - } catch (e: Throwable) { + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { _onError.invoke(e) null } @@ -115,7 +160,9 @@ 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) } } 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 ff46263d2..51c4b329d 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,5 +1,7 @@ package io.modelcontextprotocol.kotlin.sdk.server +import io.kotest.assertions.nondeterministic.eventually +import io.kotest.matchers.shouldBe import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification @@ -9,19 +11,31 @@ import io.modelcontextprotocol.kotlin.sdk.types.toJSON import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.runBlocking import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withTimeout +import kotlinx.io.Buffer +import kotlinx.io.RawSink +import kotlinx.io.RawSource import kotlinx.io.Sink import kotlinx.io.Source import kotlinx.io.asSink import kotlinx.io.asSource import kotlinx.io.buffered +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.Arguments +import org.junit.jupiter.params.provider.MethodSource import java.io.ByteArrayOutputStream +import java.io.IOException import java.io.PipedInputStream import java.io.PipedOutputStream -import kotlin.test.BeforeTest -import kotlin.test.Test -import kotlin.test.assertEquals +import java.util.stream.Stream +import kotlin.test.assertFailsWith import kotlin.test.assertFalse import kotlin.test.assertTrue +import kotlin.time.Duration.Companion.seconds + +private val TEST_TIMEOUT = 5.seconds class StdioServerTransportTest { private lateinit var input: PipedInputStream @@ -33,7 +47,7 @@ class StdioServerTransportTest { private lateinit var bufferedInput: Source private lateinit var printOutput: Sink - @BeforeTest + @BeforeEach fun setUp() { // Simulate an input stream that we can push data into using inputWriter. input = PipedInputStream() @@ -50,28 +64,35 @@ class StdioServerTransportTest { } bufferedInput = input.asSource().buffered() - printOutput = output.asSink().buffered() } + @Test + fun `should be safe to close before start`() = runTest { + val server = StdioServerTransport(bufferedInput, printOutput) + server.close() // initialized guard makes this a no-op; must not throw + } + @Test fun `should start then close cleanly`() { runBlocking { - val server = StdioServerTransport(bufferedInput, printOutput) - server.onError { error -> - throw error - } + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onError { error -> + throw error + } - var didClose = false - server.onClose { - didClose = true - } + var didClose = false + server.onClose { + didClose = true + } - server.start() - assertFalse(didClose, "Should not have closed yet") + server.start() + assertFalse(didClose, "Should not have closed yet") - server.close() - assertTrue(didClose, "Should have closed after calling close()") + server.close() + assertTrue(didClose, "Should have closed after calling close()") + } } } @@ -101,7 +122,7 @@ class StdioServerTransportTest { server.start() val received = readMessage.await() - assertEquals(message, received) + received shouldBe message } @Test @@ -135,7 +156,152 @@ class StdioServerTransportTest { server.start() finished.await() - assertEquals(messages, readMessages) + readMessages shouldBe messages + } + + // region: Exception handling + + @ParameterizedTest(name = "[{index}] input throws {0}") + @MethodSource("inputExceptions") + fun `should invoke onError when input stream throws`(exception: Exception): Unit = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(FaultyRawSource(exception).buffered(), printOutput) + val capturedError = CompletableDeferred() + server.onError { capturedError.complete(it) } + server.onMessage {} + + server.start() + + capturedError.await() shouldBe exception + server.close() + } + } + + @ParameterizedTest(name = "[{index}] output throws {0}") + @MethodSource("outputExceptions") + fun `should invoke onError when output sink throws`(exception: Exception): Unit = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, FaultyRawSink(exception).buffered()) + val capturedError = CompletableDeferred() + server.onError { capturedError.complete(it) } + server.onMessage {} + + server.start() + server.send(PingRequest().toJSON()) + + capturedError.await() shouldBe exception + server.close() + } + } + + @Test + fun `should call onClose when input EOF is reached`(): Unit = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, printOutput) + val didClose = CompletableDeferred() + server.onError { throw it } + server.onClose { didClose.complete(Unit) } + server.onMessage {} + + server.start() + inputWriter.close() // signal EOF to the reading loop + + eventually(2.seconds) { + didClose.isCompleted shouldBe true + } + } + } + + @Test + fun `should throw when starting twice`(): Unit = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onMessage {} + server.start() + try { + assertFailsWith { + server.start() + } + } finally { + server.close() + } + } + } + + // endregion + + @Test + fun `should invoke onError when message handler throws`() = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, printOutput) + val handlerException = RuntimeException("handler failed") + val capturedError = CompletableDeferred() + server.onError { capturedError.complete(it) } + server.onMessage { throw handlerException } + + inputWriter.write(serializeMessage(PingRequest().toJSON()).toByteArray()) + inputWriter.flush() + + server.start() + + capturedError.await() shouldBe handlerException + server.close() + } + } + + @Test + fun `should continue receiving valid messages after malformed JSON is skipped`() = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, printOutput) + val received = CompletableDeferred() + // ReadBuffer silently skips unparseable lines — no onError callback expected + server.onError {} + server.onMessage { received.complete(it) } + + val validMessage = PingRequest().toJSON() + inputWriter.write("not-valid-json\n".toByteArray()) + inputWriter.write(serializeMessage(validMessage).toByteArray()) + inputWriter.flush() + + server.start() + + received.await() shouldBe validMessage + server.close() + } + } + + companion object { + @JvmStatic + fun inputExceptions(): Stream = Stream.of( + Arguments.of(IOException("simulated read failure")), + Arguments.of(RuntimeException("unexpected read error")), + ) + + @JvmStatic + fun outputExceptions(): Stream = Stream.of( + Arguments.of(IOException("simulated write failure")), + Arguments.of(RuntimeException("unexpected write error")), + ) + } + + /** A [RawSource] that immediately throws [exception] on every read attempt. */ + private class FaultyRawSource(private val exception: Exception) : RawSource { + override fun readAtMostTo(sink: Buffer, byteCount: Long): Long = throw exception + override fun close() { + // noop + } + } + + /** A [RawSink] that throws [exception] on every [write] call. */ + private class FaultyRawSink(private val exception: Exception) : RawSink { + override fun write(source: Buffer, byteCount: Long): Unit = throw exception + override fun flush() { + // noop + } + + override fun close() { + // noop + } } } From 4ef9d9fa12e57be97cab12edd3f90c09d48cd01c Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Mon, 2 Mar 2026 11:13:17 +0200 Subject: [PATCH 2/7] fix: ReadBuffer skips malformed lines instead of stalling on them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../kotlin/sdk/shared/ReadBuffer.kt | 88 ++++++++++++------- .../kotlin/sdk/shared/ReadBufferTest.kt | 16 ++++ .../sdk/server/StdioServerTransportTest.kt | 1 + 3 files changed, 73 insertions(+), 32 deletions(-) diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt index d991e7fe0..feb24f6df 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt @@ -20,45 +20,69 @@ 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 + var message: JSONRPCMessage? = null + while (message == null) { + val line = readNextLine() ?: return null + if (line.isBlank()) continue - 0L -> { - buffer.skip(1) - return null - } - - 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 + @Suppress("TooGenericExceptionCaught") + try { + message = deserializeMessage(line) + } catch (e: Exception) { + logger.error(e) { "Failed to deserialize message from line: $line\nAttempting to recover..." } + message = tryRecover(line) } } - 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..." } - } + return message + } + + 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 } + } - return null + 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 + + 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() { diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt index a49ff3df6..7b4f05308 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferTest.kt +++ b/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() 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 51c4b329d..e1ad3e59c 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 @@ -260,6 +260,7 @@ class StdioServerTransportTest { val validMessage = PingRequest().toJSON() inputWriter.write("not-valid-json\n".toByteArray()) + inputWriter.write(" \t \r \n".toByteArray()) // blank line inputWriter.write(serializeMessage(validMessage).toByteArray()) inputWriter.flush() From 631862315d8f41a1c0fe02822b374d400dfe641a Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Mon, 2 Mar 2026 17:21:05 +0200 Subject: [PATCH 3/7] Address PR comments --- .../kotlin/sdk/shared/ReadBuffer.kt | 13 +++++++------ .../kotlin/sdk/server/StdioServerTransport.kt | 19 +++++++------------ .../sdk/server/StdioServerTransportTest.kt | 10 +++++----- 3 files changed, 19 insertions(+), 23 deletions(-) diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt index feb24f6df..d3fc61fd4 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ReadBuffer.kt @@ -34,20 +34,21 @@ public class ReadBuffer { * input buffer is exhausted or no valid message is found. */ public fun readMessage(): JSONRPCMessage? { - var message: JSONRPCMessage? = null - while (message == null) { + while (true) { val line = readNextLine() ?: return null if (line.isBlank()) continue @Suppress("TooGenericExceptionCaught") - try { - message = deserializeMessage(line) + val message = try { + deserializeMessage(line) } catch (e: Exception) { logger.error(e) { "Failed to deserialize message from line: $line\nAttempting to recover..." } - message = tryRecover(line) + tryRecover(line) + } + if (message != null) { + return message } } - return message } private fun readNextLine(): String? { 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 9a4eee884..a0b6ba958 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 @@ -27,6 +27,8 @@ 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. * @@ -39,11 +41,7 @@ import kotlin.coroutines.CoroutineContext @OptIn(ExperimentalAtomicApi::class) public class StdioServerTransport(private val inputStream: Source, outputStream: Sink) : AbstractTransport() { - private companion object { - private val logger = KotlinLogging.logger {} - private const val READ_BUFFER_SIZE = 8192L - } - + private val logger = KotlinLogging.logger {} private val readBuffer = ReadBuffer() private val initialized: AtomicBoolean = AtomicBoolean(false) private var readingJob: Job? = null @@ -55,7 +53,6 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: private val writeChannel = Channel(Channel.UNLIMITED) private val outputWriter = outputStream.buffered() - @Suppress("TooGenericExceptionCaught", "ThrowsCount") override suspend fun start() { if (!initialized.compareAndSet(expectedValue = false, newValue = true)) { error("StdioServerTransport already started!") @@ -71,10 +68,10 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: sendingJob = launchSendingJob() } - @Suppress("TooGenericExceptionCaught") private fun launchReadingJob(): Job { val job = scope.launch { val buf = Buffer() + @Suppress("TooGenericExceptionCaught") try { while (isActive) { val bytesRead = inputStream.readAtMostTo(buf, READ_BUFFER_SIZE) @@ -103,9 +100,9 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: return job } - @Suppress("TooGenericExceptionCaught") private fun launchProcessingJob() { scope.launch { + @Suppress("TooGenericExceptionCaught") try { for (chunk in readChannel) { readBuffer.append(chunk) @@ -119,9 +116,9 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } - @Suppress("TooGenericExceptionCaught") private fun launchSendingJob(): Job { val job = scope.launch { + @Suppress("TooGenericExceptionCaught") try { for (message in writeChannel) { val json = serializeMessage(message) @@ -144,13 +141,11 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: return job } - @Suppress("TooGenericExceptionCaught") private suspend fun processReadBuffer() { + @Suppress("TooGenericExceptionCaught") while (true) { val message = try { readBuffer.readMessage() - } catch (e: CancellationException) { - throw e } catch (e: Exception) { _onError.invoke(e) null 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 e1ad3e59c..f3ef0081b 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,8 @@ 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.shouldBe import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage @@ -30,7 +32,6 @@ import java.io.IOException import java.io.PipedInputStream import java.io.PipedOutputStream import java.util.stream.Stream -import kotlin.test.assertFailsWith import kotlin.test.assertFalse import kotlin.test.assertTrue import kotlin.time.Duration.Companion.seconds @@ -218,13 +219,12 @@ class StdioServerTransportTest { val server = StdioServerTransport(bufferedInput, printOutput) server.onMessage {} server.start() - try { - assertFailsWith { + withClue("Server should not start twice") { + shouldThrow { server.start() } - } finally { - server.close() } + server.close() } } From 2bfefa4929249634adfeed9a69f90facb18001f0 Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Mon, 2 Mar 2026 21:27:17 +0200 Subject: [PATCH 4/7] test: verify exception handling in StdioServerTransport - Add test verifying server continues processing messages despite handler exceptions. - Introduce `ParameterizedTest` for various exception types in message handler. - Ensure `onError` is not invoked for `CancellationException`. --- .../sdk/server/StdioServerTransportTest.kt | 73 ++++++++++++++++--- 1 file changed, 64 insertions(+), 9 deletions(-) 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 f3ef0081b..4a823bfa3 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 @@ -3,6 +3,7 @@ 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 import io.kotest.matchers.shouldBe import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage @@ -10,7 +11,9 @@ import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.PingRequest import io.modelcontextprotocol.kotlin.sdk.types.toJSON +import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.TimeoutCancellationException import kotlinx.coroutines.runBlocking import kotlinx.coroutines.test.runTest import kotlinx.coroutines.withTimeout @@ -34,6 +37,7 @@ import java.io.PipedOutputStream import java.util.stream.Stream import kotlin.test.assertFalse import kotlin.test.assertTrue +import kotlin.test.fail import kotlin.time.Duration.Companion.seconds private val TEST_TIMEOUT = 5.seconds @@ -228,27 +232,72 @@ class StdioServerTransportTest { } } - // endregion + @ParameterizedTest(name = "[{index}] handler throws {0}") + @MethodSource("handlerExceptions") + fun `should continue processing messages after handler throws`(exception: Exception) = runBlocking { + withTimeout(TEST_TIMEOUT) { + val server = StdioServerTransport(bufferedInput, printOutput) + val capturedErrors = mutableListOf() + val receivedMessages = mutableListOf() + val secondMessageProcessed = CompletableDeferred() + + val message1 = PingRequest().toJSON() + val message2 = InitializedNotification().toJSON() + + server.onError { capturedErrors.add(it) } + server.onMessage { message -> + if (message == message1) { + throw exception + } else { + receivedMessages.add(message) + secondMessageProcessed.complete(Unit) + } + } + + server.start() + + inputWriter.write(serializeMessage(message1)) + inputWriter.write(serializeMessage(message2)) + inputWriter.flush() + + secondMessageProcessed.await() + + capturedErrors shouldContain exception + receivedMessages shouldBe listOf(message2) + server.close() + } + } @Test - fun `should invoke onError when message handler throws`() = runBlocking { + fun `should not invoke onError for CancellationException in handler`() = runBlocking { withTimeout(TEST_TIMEOUT) { val server = StdioServerTransport(bufferedInput, printOutput) - val handlerException = RuntimeException("handler failed") val capturedError = CompletableDeferred() server.onError { capturedError.complete(it) } - server.onMessage { throw handlerException } - - inputWriter.write(serializeMessage(PingRequest().toJSON()).toByteArray()) - inputWriter.flush() + server.onMessage { throw CancellationException("cancelled") } server.start() - capturedError.await() shouldBe handlerException - server.close() + inputWriter.write(serializeMessage(PingRequest().toJSON())) + inputWriter.flush() + + // We expect onError NOT to be called. + // We wait a bit to make sure it's not called, then close. + try { + withTimeout(1.seconds) { + capturedError.await() + } + fail("Should not have captured an error for CancellationException") + } catch (_: TimeoutCancellationException) { + // Success - timeout reached without error captured + } finally { + server.close() + } } } + // endregion + @Test fun `should continue receiving valid messages after malformed JSON is skipped`() = runBlocking { withTimeout(TEST_TIMEOUT) { @@ -283,6 +332,12 @@ class StdioServerTransportTest { Arguments.of(IOException("simulated write failure")), Arguments.of(RuntimeException("unexpected write error")), ) + + @JvmStatic + fun handlerExceptions(): Stream = Stream.of( + Arguments.of(RuntimeException("handler failure")), + Arguments.of(IOException("handler IO failure")), + ) } /** A [RawSource] that immediately throws [exception] on every read attempt. */ From 5572cf6bab1560147b3f86cf3dcd87c13bc4f806 Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Mon, 2 Mar 2026 22:13:20 +0200 Subject: [PATCH 5/7] refactor(test): replace `runBlocking` with `runIntegrationTest` in StdioServerTransport tests - Introduced `runIntegrationTest` utility for consistent test execution with timeout configuration. - Adjusted test utils module in Gradle build to include new helper. - Simplified exception handling and concurrency in StdioServerTransportTest. --- kotlin-sdk-server/build.gradle.kts | 1 + .../sdk/server/StdioServerTransportTest.kt | 295 ++++++++---------- .../kotlin/test/utils/testing.kt | 29 ++ 3 files changed, 166 insertions(+), 159 deletions(-) create mode 100644 test-utils/src/jvmMain/kotlin/io/modelcontextprotocol/kotlin/test/utils/testing.kt diff --git a/kotlin-sdk-server/build.gradle.kts b/kotlin-sdk-server/build.gradle.kts index 134b36869..3a0a3985a 100644 --- a/kotlin-sdk-server/build.gradle.kts +++ b/kotlin-sdk-server/build.gradle.kts @@ -23,6 +23,7 @@ kotlin { implementation(libs.kotest.assertions.core) implementation(libs.kotest.assertions.json) implementation(libs.kotlinx.coroutines.test) + implementation(project(":test-utils")) } } 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 4a823bfa3..5fff7f5c5 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 @@ -11,11 +11,10 @@ import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.PingRequest import io.modelcontextprotocol.kotlin.sdk.types.toJSON +import io.modelcontextprotocol.kotlin.test.utils.runIntegrationTest import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.TimeoutCancellationException -import kotlinx.coroutines.runBlocking -import kotlinx.coroutines.test.runTest import kotlinx.coroutines.withTimeout import kotlinx.io.Buffer import kotlinx.io.RawSink @@ -27,21 +26,19 @@ import kotlinx.io.asSource import kotlinx.io.buffered import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test +import org.junit.jupiter.api.TestInstance import org.junit.jupiter.params.ParameterizedTest -import org.junit.jupiter.params.provider.Arguments import org.junit.jupiter.params.provider.MethodSource import java.io.ByteArrayOutputStream import java.io.IOException import java.io.PipedInputStream import java.io.PipedOutputStream -import java.util.stream.Stream import kotlin.test.assertFalse import kotlin.test.assertTrue import kotlin.test.fail import kotlin.time.Duration.Companion.seconds -private val TEST_TIMEOUT = 5.seconds - +@TestInstance(TestInstance.Lifecycle.PER_CLASS) class StdioServerTransportTest { private lateinit var input: PipedInputStream private lateinit var inputWriter: PipedOutputStream @@ -73,36 +70,32 @@ class StdioServerTransportTest { } @Test - fun `should be safe to close before start`() = runTest { + fun `should be safe to close before start`() = runIntegrationTest { val server = StdioServerTransport(bufferedInput, printOutput) server.close() // initialized guard makes this a no-op; must not throw } @Test - fun `should start then close cleanly`() { - runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, printOutput) - server.onError { error -> - throw error - } - - var didClose = false - server.onClose { - didClose = true - } - - server.start() - assertFalse(didClose, "Should not have closed yet") + fun `should start then close cleanly`() = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onError { error -> + throw error + } - server.close() - assertTrue(didClose, "Should have closed after calling close()") - } + var didClose = false + server.onClose { + didClose = true } + + server.start() + assertFalse(didClose, "Should not have closed yet") + + server.close() + assertTrue(didClose, "Should have closed after calling close()") } @Test - fun `should not read until started`() = runTest { + fun `should not read until started`() = runIntegrationTest { val server = StdioServerTransport(bufferedInput, printOutput) server.onError { error -> throw error @@ -131,7 +124,7 @@ class StdioServerTransportTest { } @Test - fun `should read multiple messages`() = runTest { + fun `should read multiple messages`() = runIntegrationTest { val server = StdioServerTransport(bufferedInput, printOutput) server.onError { error -> throw error @@ -168,178 +161,162 @@ class StdioServerTransportTest { @ParameterizedTest(name = "[{index}] input throws {0}") @MethodSource("inputExceptions") - fun `should invoke onError when input stream throws`(exception: Exception): Unit = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(FaultyRawSource(exception).buffered(), printOutput) - val capturedError = CompletableDeferred() - server.onError { capturedError.complete(it) } - server.onMessage {} + fun `should invoke onError when input stream throws`(exception: Exception): Unit = runIntegrationTest { + val server = StdioServerTransport(FaultyRawSource(exception).buffered(), printOutput) + val capturedError = CompletableDeferred() + server.onError { capturedError.complete(it) } + server.onMessage {} - server.start() + server.start() - capturedError.await() shouldBe exception - server.close() - } + capturedError.await() shouldBe exception + server.close() } @ParameterizedTest(name = "[{index}] output throws {0}") @MethodSource("outputExceptions") - fun `should invoke onError when output sink throws`(exception: Exception): Unit = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, FaultyRawSink(exception).buffered()) - val capturedError = CompletableDeferred() - server.onError { capturedError.complete(it) } - server.onMessage {} + fun `should invoke onError when output sink throws`(exception: Exception): Unit = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, FaultyRawSink(exception).buffered()) + val capturedError = CompletableDeferred() + server.onError { capturedError.complete(it) } + server.onMessage {} - server.start() - server.send(PingRequest().toJSON()) + server.start() + server.send(PingRequest().toJSON()) - capturedError.await() shouldBe exception - server.close() - } + capturedError.await() shouldBe exception + server.close() } @Test - fun `should call onClose when input EOF is reached`(): Unit = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, printOutput) - val didClose = CompletableDeferred() - server.onError { throw it } - server.onClose { didClose.complete(Unit) } - server.onMessage {} - - server.start() - inputWriter.close() // signal EOF to the reading loop - - eventually(2.seconds) { - didClose.isCompleted shouldBe true - } + fun `should call onClose when input EOF is reached`(): Unit = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + val didClose = CompletableDeferred() + server.onError { throw it } + server.onClose { didClose.complete(Unit) } + server.onMessage {} + + server.start() + inputWriter.close() // signal EOF to the reading loop + + eventually(2.seconds) { + didClose.isCompleted shouldBe true } } @Test - fun `should throw when starting twice`(): Unit = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, printOutput) - server.onMessage {} - server.start() - withClue("Server should not start twice") { - shouldThrow { - server.start() - } + fun `should throw when starting twice`(): Unit = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onMessage {} + server.start() + withClue("Server should not start twice") { + shouldThrow { + server.start() } - server.close() } + server.close() } @ParameterizedTest(name = "[{index}] handler throws {0}") @MethodSource("handlerExceptions") - fun `should continue processing messages after handler throws`(exception: Exception) = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, printOutput) - val capturedErrors = mutableListOf() - val receivedMessages = mutableListOf() - val secondMessageProcessed = CompletableDeferred() - - val message1 = PingRequest().toJSON() - val message2 = InitializedNotification().toJSON() - - server.onError { capturedErrors.add(it) } - server.onMessage { message -> - if (message == message1) { - throw exception - } else { - receivedMessages.add(message) - secondMessageProcessed.complete(Unit) - } + fun `should continue processing messages after handler throws`(exception: Exception) = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + val capturedErrors = mutableListOf() + val receivedMessages = mutableListOf() + val secondMessageProcessed = CompletableDeferred() + + val message1 = PingRequest().toJSON() + val message2 = InitializedNotification().toJSON() + + server.onError { capturedErrors.add(it) } + server.onMessage { message -> + if (message == message1) { + throw exception + } else { + receivedMessages.add(message) + secondMessageProcessed.complete(Unit) } + } - server.start() + server.start() - inputWriter.write(serializeMessage(message1)) - inputWriter.write(serializeMessage(message2)) - inputWriter.flush() + inputWriter.write(serializeMessage(message1)) + inputWriter.write(serializeMessage(message2)) + inputWriter.flush() - secondMessageProcessed.await() + secondMessageProcessed.await() - capturedErrors shouldContain exception - receivedMessages shouldBe listOf(message2) - server.close() - } + capturedErrors shouldContain exception + receivedMessages shouldBe listOf(message2) + server.close() } @Test - fun `should not invoke onError for CancellationException in handler`() = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, printOutput) - val capturedError = CompletableDeferred() - server.onError { capturedError.complete(it) } - - server.onMessage { throw CancellationException("cancelled") } - server.start() - - inputWriter.write(serializeMessage(PingRequest().toJSON())) - inputWriter.flush() - - // We expect onError NOT to be called. - // We wait a bit to make sure it's not called, then close. - try { - withTimeout(1.seconds) { - capturedError.await() - } - fail("Should not have captured an error for CancellationException") - } catch (_: TimeoutCancellationException) { - // Success - timeout reached without error captured - } finally { - server.close() + fun `should not invoke onError for CancellationException in handler`() = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + val capturedError = CompletableDeferred() + server.onError { capturedError.complete(it) } + + server.onMessage { throw CancellationException("cancelled") } + server.start() + + inputWriter.write(serializeMessage(PingRequest().toJSON())) + inputWriter.flush() + + // We expect onError NOT to be called. + // We wait a bit to make sure it's not called, then close. + try { + withTimeout(1.seconds) { + capturedError.await() } + fail("Should not have captured an error for CancellationException") + } catch (_: TimeoutCancellationException) { + // Success - timeout reached without error captured + } finally { + server.close() } } // endregion @Test - fun `should continue receiving valid messages after malformed JSON is skipped`() = runBlocking { - withTimeout(TEST_TIMEOUT) { - val server = StdioServerTransport(bufferedInput, printOutput) - val received = CompletableDeferred() - // ReadBuffer silently skips unparseable lines — no onError callback expected - server.onError {} - server.onMessage { received.complete(it) } - - val validMessage = PingRequest().toJSON() - inputWriter.write("not-valid-json\n".toByteArray()) - inputWriter.write(" \t \r \n".toByteArray()) // blank line - inputWriter.write(serializeMessage(validMessage).toByteArray()) - inputWriter.flush() - - server.start() - - received.await() shouldBe validMessage - server.close() - } - } - - companion object { - @JvmStatic - fun inputExceptions(): Stream = Stream.of( - Arguments.of(IOException("simulated read failure")), - Arguments.of(RuntimeException("unexpected read error")), - ) + fun `should continue receiving valid messages after malformed JSON is skipped`() = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + val received = CompletableDeferred() + // ReadBuffer silently skips unparseable lines — no onError callback expected + server.onError {} + server.onMessage { received.complete(it) } + + val validMessage = PingRequest().toJSON() + inputWriter.write("not-valid-json\n".toByteArray()) + inputWriter.write(" \t \r \n".toByteArray()) // blank line + inputWriter.write(serializeMessage(validMessage).toByteArray()) + inputWriter.flush() - @JvmStatic - fun outputExceptions(): Stream = Stream.of( - Arguments.of(IOException("simulated write failure")), - Arguments.of(RuntimeException("unexpected write error")), - ) + server.start() - @JvmStatic - fun handlerExceptions(): Stream = Stream.of( - Arguments.of(RuntimeException("handler failure")), - Arguments.of(IOException("handler IO failure")), - ) + received.await() shouldBe validMessage + server.close() } + @Suppress("unused") + private fun inputExceptions() = listOf( + IOException("simulated read failure"), + RuntimeException("unexpected read error"), + ) + + @Suppress("unused") + private fun outputExceptions() = listOf( + IOException("simulated write failure"), + RuntimeException("unexpected write error"), + ) + + @Suppress("unused") + private fun handlerExceptions() = listOf( + RuntimeException("handler failure"), + IOException("handler IO failure"), + ) + /** A [RawSource] that immediately throws [exception] on every read attempt. */ private class FaultyRawSource(private val exception: Exception) : RawSource { override fun readAtMostTo(sink: Buffer, byteCount: Long): Long = throw exception diff --git a/test-utils/src/jvmMain/kotlin/io/modelcontextprotocol/kotlin/test/utils/testing.kt b/test-utils/src/jvmMain/kotlin/io/modelcontextprotocol/kotlin/test/utils/testing.kt new file mode 100644 index 000000000..ee7316254 --- /dev/null +++ b/test-utils/src/jvmMain/kotlin/io/modelcontextprotocol/kotlin/test/utils/testing.kt @@ -0,0 +1,29 @@ +package io.modelcontextprotocol.kotlin.test.utils + +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeout +import kotlin.coroutines.CoroutineContext +import kotlin.time.Duration +import kotlin.time.Duration.Companion.seconds + +@Suppress("InjectDispatcher") // this is constant +private val defaultTestDispatcher: CoroutineContext = Dispatchers.IO.limitedParallelism(100) +private val defaultTestTimeout: Duration = 5.seconds + +/** + * Executes a suspendable integration test block within a coroutine scope, applying a timeout constraint. + * + * @param timeout The maximum duration allowed for the test execution. Defaults to 5 seconds. + * @param context The coroutine context within which the test block is executed. + * Defaults to [Dispatchers.IO] with 100 parallel tasks. + * @param block The suspendable test logic to be executed. + */ +public fun runIntegrationTest( + context: CoroutineContext = defaultTestDispatcher, + timeout: Duration = defaultTestTimeout, + block: suspend CoroutineScope.() -> Unit, +): Unit = runBlocking(context) { + withTimeout(timeout, block) +} From 9d47a71b261035838e562f44604ffb77f092b917 Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Mon, 2 Mar 2026 23:40:29 +0200 Subject: [PATCH 6/7] refactor: Refactored processing job management & fix KDoc in StdioServerTransport - Refactored job management to include explicit tracking for the processing job. Enhanced resource cleanup logic. - Updated constructor parameters KDoc to source/sink to match actual types. Renamed internal variable `outputWriter` to `outputSink` to reflect updated abstraction. - Added helper method `logJobCompletion` for consistent logging of job lifecycle events. --- .../kotlin/sdk/server/StdioServerTransport.kt | 50 +++++++++++++------ .../sdk/server/StdioServerTransportTest.kt | 2 +- 2 files changed, 37 insertions(+), 15 deletions(-) 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 a0b6ba958..14fc4066d 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 @@ -32,11 +32,11 @@ 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 stream used to receive data. - * @param outputStream The output stream used to send data. + * @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() { @@ -46,12 +46,13 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: 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(Channel.UNLIMITED) private val writeChannel = Channel(Channel.UNLIMITED) - private val outputWriter = outputStream.buffered() + private val outputSink = outputStream.buffered() override suspend fun start() { if (!initialized.compareAndSet(expectedValue = false, newValue = true)) { @@ -62,7 +63,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: readingJob = launchReadingJob() // Launch a coroutine to process messages from readChannel - launchProcessingJob() + processingJob = launchProcessingJob() // Launch a coroutine to handle message sending sendingJob = launchSendingJob() @@ -95,13 +96,13 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } job.invokeOnCompletion { cause -> - logger.debug(cause) { "Message reading job completed with cause: $cause" } + logJobCompletion("Message reading", cause) } return job } - private fun launchProcessingJob() { - scope.launch { + private fun launchProcessingJob(): Job { + val job = scope.launch { @Suppress("TooGenericExceptionCaught") try { for (chunk in readChannel) { @@ -114,6 +115,10 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: _onError.invoke(e) } } + job.invokeOnCompletion { cause -> + logJobCompletion("Processing", cause) + } + return job } private fun launchSendingJob(): Job { @@ -122,8 +127,8 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: try { for (message in writeChannel) { val json = serializeMessage(message) - outputWriter.writeString(json) - outputWriter.flush() + outputSink.writeString(json) + outputSink.flush() } } catch (e: CancellationException) { throw e @@ -133,7 +138,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } job.invokeOnCompletion { cause -> - logger.debug(cause) { "Message sending job completed with cause: $cause" } + logJobCompletion("Message sending", cause) if (cause is CancellationException) { readingJob?.cancel(cause) } @@ -163,6 +168,21 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } + 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 @@ -175,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() 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 5fff7f5c5..d1ef24824 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 @@ -338,6 +338,6 @@ class StdioServerTransportTest { } } -fun PipedOutputStream.write(s: String) { +private fun PipedOutputStream.write(s: String) { write(s.toByteArray()) } From 9c0ad9f1a2c60336efddaf6c7edd0ac8453a6185 Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Tue, 3 Mar 2026 13:06:20 +0200 Subject: [PATCH 7/7] refactor: replace `Exception` with `Throwable` in StdioServerTransport and tests Addressed PR comments: - Updated exception handling in `StdioServerTransport` to catch `Throwable` instead of `Exception`. - Adjusted test methods and helper functions to handle `Throwable`. --- .../kotlin/sdk/server/StdioServerTransport.kt | 11 +++-- .../sdk/server/StdioServerTransportTest.kt | 49 ++++++++++--------- 2 files changed, 32 insertions(+), 28 deletions(-) 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 14fc4066d..8fd46dd21 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 @@ -87,7 +87,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } catch (e: CancellationException) { throw e - } catch (e: Exception) { + } catch (e: Throwable) { logger.error(e) { "Error reading from stdin" } _onError.invoke(e) } finally { @@ -111,7 +111,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } catch (e: CancellationException) { throw e - } catch (e: Exception) { + } catch (e: Throwable) { _onError.invoke(e) } } @@ -132,7 +132,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } } catch (e: CancellationException) { throw e - } catch (e: Exception) { + } catch (e: Throwable) { logger.error(e) { "Error writing to stdout" } _onError.invoke(e) } @@ -151,7 +151,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: while (true) { val message = try { readBuffer.readMessage() - } catch (e: Exception) { + } catch (e: Throwable) { _onError.invoke(e) null } @@ -162,7 +162,8 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: _onMessage.invoke(message) } catch (e: CancellationException) { throw e - } catch (e: Exception) { + } catch (e: Throwable) { + logger.error(e) { "Error processing message" } _onError.invoke(e) } } 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 d1ef24824..dcbb7a4b8 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 @@ -160,23 +160,23 @@ class StdioServerTransportTest { // region: Exception handling @ParameterizedTest(name = "[{index}] input throws {0}") - @MethodSource("inputExceptions") - fun `should invoke onError when input stream throws`(exception: Exception): Unit = runIntegrationTest { - val server = StdioServerTransport(FaultyRawSource(exception).buffered(), printOutput) + @MethodSource("inputErrors") + fun `should invoke onError when input stream throws`(throwable: Throwable): Unit = runIntegrationTest { + val server = StdioServerTransport(FaultyRawSource(throwable).buffered(), printOutput) val capturedError = CompletableDeferred() server.onError { capturedError.complete(it) } server.onMessage {} server.start() - capturedError.await() shouldBe exception + capturedError.await() shouldBe throwable server.close() } @ParameterizedTest(name = "[{index}] output throws {0}") - @MethodSource("outputExceptions") - fun `should invoke onError when output sink throws`(exception: Exception): Unit = runIntegrationTest { - val server = StdioServerTransport(bufferedInput, FaultyRawSink(exception).buffered()) + @MethodSource("outputErrors") + fun `should invoke onError when output sink throws`(throwable: Throwable): Unit = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, FaultyRawSink(throwable).buffered()) val capturedError = CompletableDeferred() server.onError { capturedError.complete(it) } server.onMessage {} @@ -184,7 +184,7 @@ class StdioServerTransportTest { server.start() server.send(PingRequest().toJSON()) - capturedError.await() shouldBe exception + capturedError.await() shouldBe throwable server.close() } @@ -218,8 +218,8 @@ class StdioServerTransportTest { } @ParameterizedTest(name = "[{index}] handler throws {0}") - @MethodSource("handlerExceptions") - fun `should continue processing messages after handler throws`(exception: Exception) = runIntegrationTest { + @MethodSource("handlerErrors") + fun `should continue processing messages after handler throws`(throwable: Throwable) = runIntegrationTest { val server = StdioServerTransport(bufferedInput, printOutput) val capturedErrors = mutableListOf() val receivedMessages = mutableListOf() @@ -231,7 +231,7 @@ class StdioServerTransportTest { server.onError { capturedErrors.add(it) } server.onMessage { message -> if (message == message1) { - throw exception + throw throwable } else { receivedMessages.add(message) secondMessageProcessed.complete(Unit) @@ -246,7 +246,7 @@ class StdioServerTransportTest { secondMessageProcessed.await() - capturedErrors shouldContain exception + capturedErrors shouldContain throwable receivedMessages shouldBe listOf(message2) server.close() } @@ -300,34 +300,37 @@ class StdioServerTransportTest { } @Suppress("unused") - private fun inputExceptions() = listOf( + private fun inputErrors() = listOf( IOException("simulated read failure"), - RuntimeException("unexpected read error"), + RuntimeException("unexpected read exception"), + OutOfMemoryError("unexpected read error"), ) @Suppress("unused") - private fun outputExceptions() = listOf( + private fun outputErrors() = listOf( IOException("simulated write failure"), - RuntimeException("unexpected write error"), + RuntimeException("unexpected write exception"), + OutOfMemoryError("unexpected write error"), ) @Suppress("unused") - private fun handlerExceptions() = listOf( + private fun handlerErrors() = listOf( RuntimeException("handler failure"), IOException("handler IO failure"), + OutOfMemoryError("handler error"), ) - /** A [RawSource] that immediately throws [exception] on every read attempt. */ - private class FaultyRawSource(private val exception: Exception) : RawSource { - override fun readAtMostTo(sink: Buffer, byteCount: Long): Long = throw exception + /** A [RawSource] that immediately throws [throwable] on every read attempt. */ + private class FaultyRawSource(private val throwable: Throwable) : RawSource { + override fun readAtMostTo(sink: Buffer, byteCount: Long): Long = throw throwable override fun close() { // noop } } - /** A [RawSink] that throws [exception] on every [write] call. */ - private class FaultyRawSink(private val exception: Exception) : RawSink { - override fun write(source: Buffer, byteCount: Long): Unit = throw exception + /** A [RawSink] that throws [throwable] on every [write] call. */ + private class FaultyRawSink(private val throwable: Throwable) : RawSink { + override fun write(source: Buffer, byteCount: Long): Unit = throw throwable override fun flush() { // noop }