From a38818018ac8570973b502600938544696fa6248 Mon Sep 17 00:00:00 2001 From: Konstantin Pavlov <1517853+kpavlov@users.noreply.github.com> Date: Tue, 3 Mar 2026 16:10:25 +0200 Subject: [PATCH] fix!(server): #575 Rethrow CancellationException in StdioServerTransport, fix ReadBuffer (#571) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 - [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 --- .../kotlin/sdk/shared/ReadBuffer.kt | 91 +++--- .../kotlin/sdk/shared/ReadBufferTest.kt | 16 ++ kotlin-sdk-server/build.gradle.kts | 2 + .../kotlin/sdk/server/StdioServerTransport.kt | 93 ++++++- .../sdk/server/StdioServerTransportTest.kt | 262 ++++++++++++++++-- .../kotlin/test/utils/testing.kt | 29 ++ 6 files changed, 416 insertions(+), 77 deletions(-) create mode 100644 test-utils/src/jvmMain/kotlin/io/modelcontextprotocol/kotlin/test/utils/testing.kt 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 d991e7f..d3fc61f 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,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 - } - } - 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..." } - } + if (message != null) { + return message } } + } - return null + 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 + + 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 a49ff3d..7b4f053 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/build.gradle.kts b/kotlin-sdk-server/build.gradle.kts index 9d0bddf..3a0a398 100644 --- a/kotlin-sdk-server/build.gradle.kts +++ b/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) 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 b3ec4d6..8fd46dd 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 @@ -26,25 +27,32 @@ 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(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)) { @@ -52,11 +60,22 @@ public class StdioServerTransport(private val inputStream: Source, outputStream: } // 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() 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 ff46263..dcbb7a4 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,28 +1,44 @@ 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 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.runBlocking -import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.TimeoutCancellationException +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.api.TestInstance +import org.junit.jupiter.params.ParameterizedTest +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 kotlin.test.assertFalse import kotlin.test.assertTrue +import kotlin.test.fail +import kotlin.time.Duration.Companion.seconds +@TestInstance(TestInstance.Lifecycle.PER_CLASS) class StdioServerTransportTest { private lateinit var input: PipedInputStream private lateinit var inputWriter: PipedOutputStream @@ -33,7 +49,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,33 +66,36 @@ class StdioServerTransportTest { } bufferedInput = input.asSource().buffered() - printOutput = output.asSink().buffered() } @Test - fun `should start then close cleanly`() { - runBlocking { - 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") - - server.close() - assertTrue(didClose, "Should have closed after calling close()") - } + 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 not read until started`() = runTest { + fun `should start then close cleanly`() = runIntegrationTest { + 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") + + server.close() + assertTrue(didClose, "Should have closed after calling close()") + } + + @Test + fun `should not read until started`() = runIntegrationTest { val server = StdioServerTransport(bufferedInput, printOutput) server.onError { error -> throw error @@ -101,11 +120,11 @@ class StdioServerTransportTest { server.start() val received = readMessage.await() - assertEquals(message, received) + received shouldBe message } @Test - fun `should read multiple messages`() = runTest { + fun `should read multiple messages`() = runIntegrationTest { val server = StdioServerTransport(bufferedInput, printOutput) server.onError { error -> throw error @@ -135,10 +154,193 @@ class StdioServerTransportTest { server.start() finished.await() - assertEquals(messages, readMessages) + readMessages shouldBe messages + } + + // region: Exception handling + + @ParameterizedTest(name = "[{index}] input throws {0}") + @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 throwable + server.close() + } + + @ParameterizedTest(name = "[{index}] output throws {0}") + @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 {} + + server.start() + server.send(PingRequest().toJSON()) + + capturedError.await() shouldBe throwable + server.close() + } + + @Test + 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 = runIntegrationTest { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onMessage {} + server.start() + withClue("Server should not start twice") { + shouldThrow { + server.start() + } + } + server.close() + } + + @ParameterizedTest(name = "[{index}] handler throws {0}") + @MethodSource("handlerErrors") + fun `should continue processing messages after handler throws`(throwable: Throwable) = 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 throwable + } else { + receivedMessages.add(message) + secondMessageProcessed.complete(Unit) + } + } + + server.start() + + inputWriter.write(serializeMessage(message1)) + inputWriter.write(serializeMessage(message2)) + inputWriter.flush() + + secondMessageProcessed.await() + + capturedErrors shouldContain throwable + receivedMessages shouldBe listOf(message2) + server.close() + } + + @Test + 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`() = 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() + + server.start() + + received.await() shouldBe validMessage + server.close() + } + + @Suppress("unused") + private fun inputErrors() = listOf( + IOException("simulated read failure"), + RuntimeException("unexpected read exception"), + OutOfMemoryError("unexpected read error"), + ) + + @Suppress("unused") + private fun outputErrors() = listOf( + IOException("simulated write failure"), + RuntimeException("unexpected write exception"), + OutOfMemoryError("unexpected write error"), + ) + + @Suppress("unused") + private fun handlerErrors() = listOf( + RuntimeException("handler failure"), + IOException("handler IO failure"), + OutOfMemoryError("handler error"), + ) + + /** 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 [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 + } + + override fun close() { + // noop + } } } -fun PipedOutputStream.write(s: String) { +private fun PipedOutputStream.write(s: String) { write(s.toByteArray()) } 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 0000000..ee73162 --- /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) +}