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..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 @@ -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() { 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/build.gradle.kts b/kotlin-sdk-server/build.gradle.kts index 9d0bddff8..3a0a3985a 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 b3ec4d650..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 @@ -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 ff46263d2..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 @@ -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 - } + 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 + } - server.start() - assertFalse(didClose, "Should not have closed yet") + @Test + 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 @@ -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 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) +}