fix(server): #573, #574: add customizable builder to StdioServerTransport

Scope lifecycle is fixed (#574 resolved), processing dispatcher defaults to `Dispatchers.Default` (#573 resolved), scope context is clean (no spurious dispatcher stacking), and @volatile on the three job vars addresses the visibility race.

- Introduce a `Configuration` class for `StdioServerTransport` to improve API flexibility and readability.
- Updated transport initialization to use a builder block for configuring parameters such as I/O streams, buffer sizes, dispatchers, and parent coroutine scope.

Test: Added integration test validating the builder functionality.

Author: kpavlov

kotlin-sdk-server/api/kotlin-sdk-server.api

@@ -186,12 +186,35 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/SseServerTransport
 }
 
 public final class io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport {
+	public fun <init> (Lkotlin/jvm/functions/Function1;)V
 	public fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;)V
 	public fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun start (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
+public final class io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport$Configuration {
+	public fun <init> ()V
+	public final fun getCoroutineScope ()Lkotlinx/coroutines/CoroutineScope;
+	public final fun getProcessingJobDispatcher ()Lkotlinx/coroutines/CoroutineDispatcher;
+	public final fun getReadBufferSize ()J
+	public final fun getReadChannelBufferSize ()I
+	public final fun getReadingJobDispatcher ()Lkotlinx/coroutines/CoroutineDispatcher;
+	public final fun getSink ()Lkotlinx/io/Sink;
+	public final fun getSource ()Lkotlinx/io/Source;
+	public final fun getWriteChannelBufferSize ()I
+	public final fun getWritingJobDispatcher ()Lkotlinx/coroutines/CoroutineDispatcher;
+	public final fun setCoroutineScope (Lkotlinx/coroutines/CoroutineScope;)V
+	public final fun setProcessingJobDispatcher (Lkotlinx/coroutines/CoroutineDispatcher;)V
+	public final fun setReadBufferSize (J)V
+	public final fun setReadChannelBufferSize (I)V
+	public final fun setReadingJobDispatcher (Lkotlinx/coroutines/CoroutineDispatcher;)V
+	public final fun setSink (Lkotlinx/io/Sink;)V
+	public final fun setSource (Lkotlinx/io/Source;)V
+	public final fun setWriteChannelBufferSize (I)V
+	public final fun setWritingJobDispatcher (Lkotlinx/coroutines/CoroutineDispatcher;)V
+}
+
 public final class io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport {
 	public static final field STANDALONE_SSE_STREAM_ID Ljava/lang/String;
 	public fun <init> ()V

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

@@ -8,7 +8,9 @@ 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.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.SupervisorJob
@@ -23,9 +25,9 @@ import kotlinx.io.Source
 import kotlinx.io.buffered
 import kotlinx.io.readByteArray
 import kotlinx.io.writeString
+import kotlin.concurrent.Volatile
 import kotlin.concurrent.atomics.AtomicBoolean
 import kotlin.concurrent.atomics.ExperimentalAtomicApi
-import kotlin.coroutines.CoroutineContext
 
 private const val READ_BUFFER_SIZE = 8192L
 
@@ -34,25 +36,91 @@ private const val READ_BUFFER_SIZE = 8192L
  *
  * 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.
+ * @constructor Initializes the transport using the provided block for configuration.
+ * The configuration includes specifying the input and output streams, buffer sizes,
+ * and dispatchers for I/O and processing tasks.
  */
 @OptIn(ExperimentalAtomicApi::class)
-public class StdioServerTransport(private val inputStream: Source, outputStream: Sink) : AbstractTransport() {
+public class StdioServerTransport(block: Configuration.() -> Unit) : AbstractTransport() {
+
+    /**
+     * Configuration for [StdioServerTransport].
+     *
+     * @property source The input [Source] used to receive data.
+     * @property sink The output [Sink] used to send data.
+     * @property readBufferSize The buffer size for the read channel.
+     * @property readingJobDispatcher The [CoroutineDispatcher] used for reading jobs.
+     *      Defaults to [IODispatcher].
+     * @property writingJobDispatcher The [CoroutineDispatcher] used for writing jobs.
+     *      Defaults to [IODispatcher].
+     * @property processingJobDispatcher The [CoroutineDispatcher] used for processing jobs.
+     *      Defaults to [Dispatchers.Default].
+     * @property readChannelBufferSize The buffer size for the read channel.
+     * @property writeChannelBufferSize The buffer size for the write channel.
+     * @property coroutineScope The [CoroutineScope] used for managing coroutines.
+     */
+    @Suppress("LongParameterList")
+    public class Configuration internal constructor(
+        public var source: Source? = null,
+        public var sink: Sink? = null,
+        public var readBufferSize: Long = READ_BUFFER_SIZE,
+        public var readingJobDispatcher: CoroutineDispatcher = IODispatcher,
+        public var writingJobDispatcher: CoroutineDispatcher = IODispatcher,
+        public var processingJobDispatcher: CoroutineDispatcher = Dispatchers.Default,
+        public var readChannelBufferSize: Int = Channel.UNLIMITED,
+        public var writeChannelBufferSize: Int = Channel.UNLIMITED,
+        public var coroutineScope: CoroutineScope? = null,
+    )
+
+    private val source: Source
+    private val sink: Sink
+    private val processingJobDispatcher: CoroutineDispatcher
+    private val readingJobDispatcher: CoroutineDispatcher
+    private val writingJobDispatcher: CoroutineDispatcher
+    private val scope: CoroutineScope
+    private val readBufferSize: Long
+    private val readChannel: Channel<ByteArray>
+    private val writeChannel: Channel<JSONRPCMessage>
+
+    init {
+        val config = Configuration().apply(block)
+        val input = requireNotNull(config.source) { "source is required" }
+        val output = requireNotNull(config.sink) { "sink is required" }
+        require(config.readBufferSize > 0) { "readBufferSize must be > 0" }
+
+        source = input
+        processingJobDispatcher = config.processingJobDispatcher
+        readingJobDispatcher = config.readingJobDispatcher
+        writingJobDispatcher = config.writingJobDispatcher
+        val parentJob = config.coroutineScope?.coroutineContext?.get(Job)
+        scope = CoroutineScope(SupervisorJob(parentJob))
+        readBufferSize = config.readBufferSize
+        readChannel = Channel(config.readChannelBufferSize)
+        writeChannel = Channel(config.writeChannelBufferSize)
+        sink = output.buffered()
+    }
+
+    /**
+     * Creates a new instance of [StdioServerTransport]
+     * with the given [inputStream] [Source] and [outputStream] [Sink].
+     */
+    public constructor(inputStream: Source, outputStream: Sink) : this({
+        source = inputStream
+        sink = outputStream
+    })
 
     private val logger = KotlinLogging.logger {}
     private val readBuffer = ReadBuffer()
     private val initialized: AtomicBoolean = AtomicBoolean(false)
+
+    @Volatile
     private var readingJob: Job? = null
+
+    @Volatile
     private var sendingJob: Job? = null
-    private var processingJob: Job? = null
 
-    private val coroutineContext: CoroutineContext = IODispatcher + SupervisorJob()
-    private val scope = CoroutineScope(coroutineContext)
-    private val readChannel = Channel<ByteArray>(Channel.UNLIMITED)
-    private val writeChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
-    private val outputSink = outputStream.buffered()
+    @Volatile
+    private var processingJob: Job? = null
 
     override suspend fun start() {
         if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
@@ -69,81 +137,71 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
         sendingJob = launchSendingJob()
     }
 
-    private fun launchReadingJob(): Job {
-        val job = scope.launch {
-            val buf = Buffer()
-            @Suppress("TooGenericExceptionCaught")
-            try {
-                while (isActive) {
-                    val bytesRead = inputStream.readAtMostTo(buf, READ_BUFFER_SIZE)
-                    if (bytesRead == -1L) {
-                        // EOF reached
-                        break
-                    }
-                    if (bytesRead > 0) {
-                        val chunk = buf.readByteArray()
-                        readChannel.send(chunk)
-                    }
+    private fun launchReadingJob(): Job = scope.launch(readingJobDispatcher) {
+        val buf = Buffer()
+        @Suppress("TooGenericExceptionCaught")
+        try {
+            while (isActive) {
+                val bytesRead = source.readAtMostTo(buf, readBufferSize)
+                if (bytesRead == -1L) {
+                    // EOF reached
+                    break
+                }
+                if (bytesRead > 0) {
+                    val chunk = buf.readByteArray()
+                    readChannel.send(chunk)
                 }
-            } catch (e: CancellationException) {
-                throw e
-            } catch (e: Throwable) {
-                logger.error(e) { "Error reading from stdin" }
-                _onError.invoke(e)
-            } finally {
-                // Reached EOF or error, close connection
-                close()
             }
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Throwable) {
+            logger.error(e) { "Error reading from stdin" }
+            _onError.invoke(e)
+        } finally {
+            // Reached EOF or error, close connection
+            close()
         }
-        job.invokeOnCompletion { cause ->
-            logJobCompletion("Message reading", cause)
-        }
-        return job
+    }.apply {
+        invokeOnCompletion { logJobCompletion("Message reading", it) }
     }
 
-    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)
+    private fun launchProcessingJob(): Job = scope.launch(processingJobDispatcher) {
+        @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
+    }.apply {
+        invokeOnCompletion { logJobCompletion("Processing", it) }
     }
 
-    private fun launchSendingJob(): Job {
-        val job = scope.launch {
-            @Suppress("TooGenericExceptionCaught")
-            try {
-                for (message in writeChannel) {
-                    val json = serializeMessage(message)
-                    outputSink.writeString(json)
-                    outputSink.flush()
-                }
-            } catch (e: CancellationException) {
-                throw e
-            } catch (e: Throwable) {
-                logger.error(e) { "Error writing to stdout" }
-                _onError.invoke(e)
+    private fun launchSendingJob(): Job = scope.launch(writingJobDispatcher) {
+        @Suppress("TooGenericExceptionCaught")
+        try {
+            for (message in writeChannel) {
+                val json = serializeMessage(message)
+                sink.writeString(json)
+                sink.flush()
             }
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Throwable) {
+            logger.error(e) { "Error writing to stdout" }
+            _onError.invoke(e)
         }
-        job.invokeOnCompletion { cause ->
+    }.apply {
+        invokeOnCompletion { cause ->
             logJobCompletion("Message sending", cause)
             if (cause is CancellationException) {
                 readingJob?.cancel(cause)
             }
         }
-        return job
     }
 
     private suspend fun processReadBuffer() {
@@ -192,7 +250,7 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
             sendingJob?.cancelAndJoin()
 
             runCatching {
-                inputStream.close()
+                source.close()
             }.onFailure { logger.warn(it) { "Failed to close stdin" } }
 
             readingJob?.cancel()
@@ -201,10 +259,9 @@ public class StdioServerTransport(private val inputStream: Source, outputStream:
             processingJob?.cancelAndJoin()
 
             readBuffer.clear()
-
             runCatching {
-                outputSink.flush()
-                outputSink.close()
+                sink.flush()
+                sink.close()
             }.onFailure { logger.warn(it) { "Failed to close stdout" } }
 
             invokeOnCloseCallback()

kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransportTest.kt

@@ -5,6 +5,8 @@ import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.assertions.withClue
 import io.kotest.matchers.collections.shouldContain
 import io.kotest.matchers.shouldBe
+import io.ktor.client.utils.clientDispatcher
+import io.ktor.utils.io.InternalAPI
 import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
 import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
 import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification
@@ -14,7 +16,10 @@ 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.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.withTimeout
 import kotlinx.io.Buffer
 import kotlinx.io.RawSink
@@ -69,6 +74,40 @@ class StdioServerTransportTest {
         printOutput = output.asSink().buffered()
     }
 
+    @OptIn(InternalAPI::class)
+    @Test
+    fun `should construct with builder`() = runIntegrationTest {
+        val received = CompletableDeferred<JSONRPCMessage>()
+
+        val ioDispatcher = Dispatchers.IO.limitedParallelism(4)
+
+        // Set every configuration parameter explicitly with non-default values,
+        // then verify a message round-trips correctly.
+        val server = StdioServerTransport {
+            source = bufferedInput
+            sink = printOutput
+            readBufferSize = 16L // non-default: smaller read chunk
+            readingJobDispatcher = ioDispatcher // non-default: limited parallelism
+            writingJobDispatcher = ioDispatcher // non-default: limited parallelism
+            processingJobDispatcher = Dispatchers.clientDispatcher(2, "Worker") // non-default
+            readChannelBufferSize = Channel.BUFFERED // non-default: bounded
+            writeChannelBufferSize = Channel.BUFFERED // non-default: bounded
+            coroutineScope = CoroutineScope(Dispatchers.Default) // non-default: parent scope provided
+        }
+        server.onError { throw it }
+        server.onMessage { received.complete(it) }
+
+        server.start()
+
+        val message = PingRequest().toJSON()
+        inputWriter.write(serializeMessage(message))
+        inputWriter.flush()
+
+        received.await() shouldBe message
+
+        server.close()
+    }
+
     @Test
     fun `should be safe to close before start`() = runIntegrationTest {
         val server = StdioServerTransport(bufferedInput, printOutput)