fix!: harden lifecycle, dispatchers, back-pressure (#778)

fixes #573 
fixes #574 
fixes #708 

- Handlers ran on `Dispatchers.IO` → now `Dispatchers.Default`,
configurable via `handlerDispatcher`.
- Reader/writer dispatcher exposed as `ioDispatcher` (default
`IODispatcher.limitedParallelism(2)`).
- Optional caller-supplied `CoroutineScope` so the pipeline isn't
detached from the caller's tree.
- `CancellationException` consistently rethrown across all pumps.
- Unbounded channels → bounded with graceful drain on close (in-flight
outbound flushed before `onClose`).
- Input `Source` now closed on natural EOF / reader error (was leaking
for non-`System.in` sources).
- Race between `start()` and `close()` fixed via 3-state CAS +
`setupComplete` barrier replacing `AtomicBoolean`.

## How Has This Been Tested?
unit and integration tests 

## Breaking Changes
- Handlers default to `Dispatchers.Default` (was `Dispatchers.IO`).
Override with `handlerDispatcher = Dispatchers.IO` if your handlers
perform blocking I/O.
- `send()` suspends under back-pressure instead of buffering unbounded.
- `StdioServerTransport(input, output)` ctor is `@Deprecated` (WARNING)
in favor of `StdioServerTransport(input, output) { ... }`. Still
binary-compatible.

## Types of changes
- [x] Bug fix (non-breaking change which fixes an issue)
- [x] 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
- [x] I have added appropriate error handling
- [x] I have added or updated documentation as needed

.gitignore

@@ -56,7 +56,7 @@ node_modules
 dist
 
 ### SWE agents ###
-.claude/settings.local.json
+.claude/
 .junie/
 
 ### Conformance test results ###

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

@@ -219,11 +219,22 @@ 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> (Lkotlinx/io/Source;Lkotlinx/io/Sink;)V
+	public fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;Lkotlin/jvm/functions/Function1;)V
+	public synthetic fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)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$Builder {
+	public final fun getHandlerDispatcher ()Lkotlinx/coroutines/CoroutineDispatcher;
+	public final fun getIoDispatcher ()Lkotlinx/coroutines/CoroutineDispatcher;
+	public final fun getScope ()Lkotlinx/coroutines/CoroutineScope;
+	public final fun setHandlerDispatcher (Lkotlinx/coroutines/CoroutineDispatcher;)V
+	public final fun setIoDispatcher (Lkotlinx/coroutines/CoroutineDispatcher;)V
+	public final fun setScope (Lkotlinx/coroutines/CoroutineScope;)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

@@ -7,13 +7,23 @@ 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 io.modelcontextprotocol.kotlin.sdk.types.McpDsl
+import io.modelcontextprotocol.kotlin.sdk.types.McpException
+import io.modelcontextprotocol.kotlin.sdk.types.RPCError
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.ClosedSendChannelException
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -23,191 +33,345 @@ import kotlinx.io.Source
 import kotlinx.io.buffered
 import kotlinx.io.readByteArray
 import kotlinx.io.writeString
-import kotlin.concurrent.atomics.AtomicBoolean
+import kotlin.concurrent.atomics.AtomicReference
 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.
+ * Server-side MCP transport that exchanges JSON-RPC messages over arbitrary byte streams.
  *
- * Reads from input [Source] and writes to output [Sink].
+ * Reads framed messages from the supplied [Source] and writes framed messages to the supplied
+ * [Sink]. Three internal coroutines drive the pipeline:
  *
- * @constructor Creates a new instance of [StdioServerTransport].
+ * - **reader** — pulls bytes from the input source into the parsing buffer; runs on
- * @param inputStream The input [Source] used to receive data.
+ *   [Builder.ioDispatcher].
- * @param outputStream The output [Sink] used to send data.
+ * - **processor** — parses messages out of the buffer and invokes the registered message handler;
+ *   runs on [Builder.handlerDispatcher] (defaults to [Dispatchers.Default]) so blocking handler
+ *   code does not starve the I/O pool.
+ * - **writer** — serialises outbound messages and flushes them to the output sink; runs on
+ *   [Builder.ioDispatcher].
+ *
+ * Both internal channels are bounded, so a slow handler or slow output naturally back-pressures
+ * the upstream producer — [send] suspends when the outbound channel is full.
+ *
+ * Both explicit [close] and a natural EOF from the input perform a graceful drain: in-flight
+ * outbound messages are flushed before `onClose` fires, and the input source and output sink are
+ * released.
+ *
+ * Example:
+ * ```kotlin
+ * val transport = StdioServerTransport(
+ *     input = System.`in`.asSource().buffered(),
+ *     output = System.out.asSink().buffered(),
+ * ) {
+ *     scope = myScope
+ * }
+ * transport.start()
+ * ```
  */
 @OptIn(ExperimentalAtomicApi::class)
-public class StdioServerTransport(private val inputStream: Source, outputStream: Sink) : AbstractTransport() {
+public class StdioServerTransport private constructor(
+    private val input: Source,
+    output: Sink,
+    private val scope: CoroutineScope?,
+    private val handlerDispatcher: CoroutineDispatcher,
+    private val ioDispatcher: CoroutineDispatcher,
+) : AbstractTransport() {
+
+    /**
+     * Creates a [StdioServerTransport] reading from [input] and writing to [output], with
+     * optional configuration applied through a [Builder] block.
+     *
+     * @param input source the transport reads JSON-RPC messages from
+     * @param output sink the transport writes JSON-RPC messages to
+     * @param block configuration applied to the underlying [Builder]
+     */
+    public constructor(
+        input: Source,
+        output: Sink,
+        block: Builder.() -> Unit = {},
+    ) : this(Builder(input, output).apply(block))
+
+    private constructor(builder: Builder) : this(
+        input = builder.input,
+        output = builder.output,
+        scope = builder.scope,
+        handlerDispatcher = builder.handlerDispatcher,
+        ioDispatcher = builder.ioDispatcher,
+    )
+
+    /**
+     * Creates a [StdioServerTransport] from the given input source and output sink with default
+     * configuration. Retained for binary compatibility; prefer the [Builder]-based constructor.
+     *
+     * @param inputStream source the transport reads JSON-RPC messages from
+     * @param outputStream sink the transport writes JSON-RPC messages to
+     */
+    @Deprecated(
+        message = "Use StdioServerTransport(input, output) { ... } instead.",
+        replaceWith = ReplaceWith("StdioServerTransport(input = inputStream, output = outputStream)"),
+        level = DeprecationLevel.WARNING,
+    )
+    public constructor(inputStream: Source, outputStream: Sink) : this(
+        input = inputStream,
+        output = outputStream,
+        scope = null,
+        handlerDispatcher = Dispatchers.Default,
+        ioDispatcher = IODispatcher.limitedParallelism(2),
+    )
 
     private val logger = KotlinLogging.logger {}
     private val readBuffer = ReadBuffer()
-    private val initialized: AtomicBoolean = AtomicBoolean(false)
+    private val readChannel = Channel<ByteArray>(Channel.BUFFERED)
-    private var readingJob: Job? = null
+    private val writeChannel = Channel<JSONRPCMessage>(Channel.BUFFERED)
-    private var sendingJob: Job? = null
+    private val outputSink = output.buffered()
-    private var processingJob: Job? = null
 
-    private val coroutineContext: CoroutineContext = IODispatcher + SupervisorJob()
+    private enum class State { New, Operational, Stopped }
-    private val scope = CoroutineScope(coroutineContext)
-    private val readChannel = Channel<ByteArray>(Channel.UNLIMITED)
-    private val writeChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
-    private val outputSink = outputStream.buffered()
 
+    private val state: AtomicReference<State> = AtomicReference(State.New)
+
+    private var readerJob: Job? = null
+    private var processorJob: Job? = null
+    private var writerJob: Job? = null
+
+    private var effectiveScope: CoroutineScope? = null
+    private val ownsScope: Boolean = (scope == null)
+
+    private val setupComplete = CompletableDeferred<Unit>()
+
+    /**
+     * Starts the reader, processor, and writer coroutines. Must be called exactly once before
+     * messages can be exchanged; subsequent calls throw.
+     */
     override suspend fun start() {
-        if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
+        if (!state.compareAndSet(State.New, State.Operational)) {
-            error("StdioServerTransport already started!")
+            when (state.load()) {
+                State.Stopped -> error("StdioServerTransport is already closed!")
+                else -> error("StdioServerTransport already started!")
+            }
         }
 
-        // Launch a coroutine to read from stdin
+        try {
-        readingJob = launchReadingJob()
+            val resolvedScope = scope ?: CoroutineScope(
+                currentCoroutineContext() + IODispatcher + SupervisorJob(),
+            )
+            effectiveScope = resolvedScope
 
-        // Launch a coroutine to process messages from readChannel
+            readerJob = resolvedScope.launch(
-        processingJob = launchProcessingJob()
+                ioDispatcher + CoroutineName("StdioServerTransport.reader"),
-
+            ) { readerPump() }
-        // Launch a coroutine to handle message sending
+            processorJob = resolvedScope.launch(
-        sendingJob = launchSendingJob()
+                handlerDispatcher + CoroutineName("StdioServerTransport.processor"),
+                start = CoroutineStart.UNDISPATCHED,
+            ) { processorPump() }
+            writerJob = resolvedScope.launch(
+                ioDispatcher + CoroutineName("StdioServerTransport.writer"),
+                start = CoroutineStart.UNDISPATCHED,
+            ) { writerPump() }
+        } finally {
+            setupComplete.complete(Unit)
+        }
     }
 
-    private fun launchReadingJob(): Job {
+    private suspend fun readerPump() {
-        val job = scope.launch {
+        val buf = Buffer()
-            val buf = Buffer()
+        try {
-            try {
+            while (currentCoroutineContext().isActive) {
-                while (isActive) {
+                val bytesRead = input.readAtMostTo(buf, READ_BUFFER_SIZE)
-                    val bytesRead = inputStream.readAtMostTo(buf, READ_BUFFER_SIZE)
+                if (bytesRead == -1L) break
-                    if (bytesRead == -1L) {
+                if (bytesRead > 0) {
-                        // EOF reached
+                    val chunk = buf.readByteArray()
-                        break
+                    readChannel.send(chunk)
-                    }
+                }
-                    if (bytesRead > 0) {
+            }
-                        val chunk = buf.readByteArray()
+        } catch (e: CancellationException) {
-                        readChannel.send(chunk)
+            throw e
+        } catch (e: Throwable) {
+            if (state.load() == State.Stopped) {
+                logger.debug(e) { "Reader interrupted by close()" }
+            } else {
+                logger.error(e) { "Error reading from input source" }
+                _onError(e)
+            }
+        } finally {
+            readChannel.close()
+        }
+    }
+
+    private suspend fun processorPump() {
+        try {
+            for (chunk in readChannel) {
+                readBuffer.append(chunk)
+                while (true) {
+                    val message = try {
+                        readBuffer.readMessage()
+                    } catch (e: CancellationException) {
+                        throw e
+                    } catch (e: Throwable) {
+                        _onError(e)
+                        null
+                    } ?: break
+                    try {
+                        _onMessage(message)
+                    } catch (e: CancellationException) {
+                        throw e
+                    } catch (e: Throwable) {
+                        logger.error(e) { "Error processing message" }
+                        _onError(e)
                     }
                 }
-            } 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
-    }
-
-    private fun launchProcessingJob(): Job {
-        val job = scope.launch {
-            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
-    }
-
-    private fun launchSendingJob(): Job {
-        val job = scope.launch {
-            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)
-            }
-        }
-        job.invokeOnCompletion { cause ->
-            logJobCompletion("Message sending", cause)
-            if (cause is CancellationException) {
-                readingJob?.cancel(cause)
-            }
-        }
-        return job
-    }
-
-    private suspend fun processReadBuffer() {
-        while (true) {
-            val message = try {
-                readBuffer.readMessage()
-            } catch (e: Throwable) {
-                _onError.invoke(e)
-                null
-            }
-
-            if (message == null) break
-            // 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)
             }
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Throwable) {
+            _onError(e)
+        } finally {
+            writeChannel.close()
         }
     }
 
-    private fun logJobCompletion(jobName: String, cause: Throwable?) {
+    private suspend fun writerPump() {
-        when (cause) {
+        try {
-            is CancellationException -> {
+            for (message in writeChannel) {
-            }
+                val json = serializeMessage(message)
-
+                outputSink.writeString(json)
-            null -> {
+                outputSink.flush()
-                logger.debug { "$jobName job completed" }
-            }
-
-            else -> {
-                logger.debug(cause) { "$jobName job completed exceptionally" }
             }
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Throwable) {
+            logger.error(e) { "Error writing to output sink" }
+            _onError(e)
+        } finally {
+            transitionToStoppedNaturally()
         }
     }
 
+    /**
+     * Closes the transport. When called after [start], waits for in-flight outbound messages to be
+     * flushed, releases the input source and output sink, cancels the internal scope when the
+     * transport owns it, and invokes `onClose`. When called before [start], transitions directly
+     * to the closed state without releasing the input/output or invoking `onClose` — the caller
+     * remains responsible for resources that were never handed off to a running transport. Safe to
+     * call multiple times and safe to race with [start].
+     */
     override suspend fun close() {
-        if (!initialized.compareAndSet(expectedValue = true, newValue = false)) return
+        var previous: State
+        do {
+            previous = state.load()
+        } while (previous != State.Stopped && !state.compareAndSet(previous, State.Stopped))
+
+        if (previous == State.New) {
+            setupComplete.complete(Unit)
+            return
+        }
 
         withContext(NonCancellable) {
-            writeChannel.close()
+            setupComplete.await()
-            sendingJob?.cancelAndJoin()
 
-            runCatching {
+            if (previous == State.Stopped) {
-                inputStream.close()
+                writerJob?.join()
-            }.onFailure { logger.warn(it) { "Failed to close stdin" } }
+                return@withContext
+            }
 
-            readingJob?.cancel()
+            runCatching { input.close() }
+                .onFailure { logger.warn(it) { "Failed to close input source" } }
+            readerJob?.cancel()
             readChannel.close()
-
+            processorJob?.cancelAndJoin()
-            processingJob?.cancelAndJoin()
+            writeChannel.close()
-
+            writerJob?.join()
+            runCatching { outputSink.close() }
+                .onFailure { logger.warn(it) { "Failed to close output sink" } }
             readBuffer.clear()
-
+            if (ownsScope) {
-            runCatching {
+                effectiveScope?.coroutineContext?.get(Job)?.cancel()
-                outputSink.flush()
+            }
-                outputSink.close()
-            }.onFailure { logger.warn(it) { "Failed to close stdout" } }
-
             invokeOnCloseCallback()
         }
     }
 
+    /**
+     * Queues [message] for the writer coroutine. Suspends when the outbound channel is full,
+     * applying back-pressure to the caller. Throws [McpException] with
+     * [RPCError.ErrorCode.CONNECTION_CLOSED] if the transport has not been started or has
+     * already closed.
+     *
+     * @param message JSON-RPC message to send
+     * @param options transport-specific send options; ignored by this transport
+     */
     override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
-        writeChannel.send(message)
+        when (state.load()) {
+            State.New -> throw McpException(
+                code = RPCError.ErrorCode.CONNECTION_CLOSED,
+                message = "Transport is not started",
+            )
+
+            State.Stopped -> throw McpException(
+                code = RPCError.ErrorCode.CONNECTION_CLOSED,
+                message = "Transport is closed",
+            )
+
+            State.Operational -> Unit
+        }
+        try {
+            writeChannel.send(message)
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: ClosedSendChannelException) {
+            throw McpException(
+                code = RPCError.ErrorCode.CONNECTION_CLOSED,
+                message = "Transport is closed",
+                cause = e,
+            )
+        }
+    }
+
+    private fun transitionToStoppedNaturally() {
+        if (!state.compareAndSet(State.Operational, State.Stopped)) return
+        runCatching { input.close() }
+            .onFailure { logger.warn(it) { "Failed to close input source" } }
+        runCatching { outputSink.close() }
+            .onFailure { logger.warn(it) { "Failed to close output sink" } }
+        readBuffer.clear()
+        if (ownsScope) {
+            // Non-null in practice: writer launches after effectiveScope is published.
+            effectiveScope?.coroutineContext?.get(Job)?.cancel()
+        }
+        invokeOnCloseCallback()
+    }
+
+    /**
+     * Configuration builder for [StdioServerTransport]. Used via the
+     * `StdioServerTransport(input, output) { ... }` factory; the I/O endpoints are supplied
+     * positionally, while [scope], [handlerDispatcher], and [ioDispatcher] are configurable
+     * inside the block.
+     *
+     * Example:
+     * ```kotlin
+     * StdioServerTransport(input, output) {
+     *     scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
+     *     handlerDispatcher = myHandlerDispatcher
+     *     ioDispatcher = Dispatchers.IO
+     * }
+     * ```
+     *
+     * @property scope Optional caller-supplied [CoroutineScope]. When non-null, the transport's
+     *   pipeline coroutines run as children of this scope and the transport does not tear it
+     *   down on [close]. When `null`, the transport creates and owns an internal scope.
+     * @property handlerDispatcher Dispatcher for invoking the registered message handler.
+     *   Defaults to [Dispatchers.Default].
+     * @property ioDispatcher Dispatcher for the reader and writer coroutines. Must allow at
+     *   least two threads to run concurrently so the reader and writer don't block each other.
+     *   The default is a two-thread view of the platform I/O dispatcher
+     *   (`IODispatcher.limitedParallelism(2)`); pass a different value to share or isolate I/O
+     *   threads with the rest of your application.
+     */
+    @McpDsl
+    public class Builder internal constructor(internal val input: Source, internal val output: Sink) {
+        public var scope: CoroutineScope? = null
+        public var handlerDispatcher: CoroutineDispatcher = Dispatchers.Default
+        public var ioDispatcher: CoroutineDispatcher = IODispatcher.limitedParallelism(2)
     }
 }

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

@@ -5,17 +5,29 @@ import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.assertions.withClue
 import io.kotest.matchers.collections.shouldContain
 import io.kotest.matchers.shouldBe
+import io.kotest.matchers.string.shouldContain
 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.McpException
 import io.modelcontextprotocol.kotlin.sdk.types.PingRequest
+import io.modelcontextprotocol.kotlin.sdk.types.RPCError
 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.SupervisorJob
 import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.withTimeout
+import kotlinx.coroutines.withTimeoutOrNull
 import kotlinx.io.Buffer
 import kotlinx.io.RawSink
 import kotlinx.io.RawSource
@@ -33,9 +45,12 @@ import java.io.ByteArrayOutputStream
 import java.io.IOException
 import java.io.PipedInputStream
 import java.io.PipedOutputStream
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.Executors
 import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 import kotlin.test.fail
+import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds
 
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
@@ -71,13 +86,13 @@ class StdioServerTransportTest {
 
     @Test
     fun `should be safe to close before start`() = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         server.close() // initialized guard makes this a no-op; must not throw
     }
 
     @Test
     fun `should start then close cleanly`() = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         server.onError { error ->
             throw error
         }
@@ -96,7 +111,7 @@ class StdioServerTransportTest {
 
     @Test
     fun `should not read until started`() = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         server.onError { error ->
             throw error
         }
@@ -125,7 +140,7 @@ class StdioServerTransportTest {
 
     @Test
     fun `should read multiple messages`() = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         server.onError { error ->
             throw error
         }
@@ -162,7 +177,7 @@ class StdioServerTransportTest {
     @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 server = StdioServerTransport(input = FaultyRawSource(throwable).buffered(), output = printOutput)
         val capturedError = CompletableDeferred<Throwable>()
         server.onError { capturedError.complete(it) }
         server.onMessage {}
@@ -176,7 +191,7 @@ class StdioServerTransportTest {
     @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 server = StdioServerTransport(input = bufferedInput, output = FaultyRawSink(throwable).buffered())
         val capturedError = CompletableDeferred<Throwable>()
         server.onError { capturedError.complete(it) }
         server.onMessage {}
@@ -190,7 +205,7 @@ class StdioServerTransportTest {
 
     @Test
     fun `should call onClose when input EOF is reached`(): Unit = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         val didClose = CompletableDeferred<Unit>()
         server.onError { throw it }
         server.onClose { didClose.complete(Unit) }
@@ -206,7 +221,7 @@ class StdioServerTransportTest {
 
     @Test
     fun `should throw when starting twice`(): Unit = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         server.onMessage {}
         server.start()
         withClue("Server should not start twice") {
@@ -217,10 +232,24 @@ class StdioServerTransportTest {
         server.close()
     }
 
+    @Test
+    fun `should report 'closed' error message when starting after close`(): Unit = runIntegrationTest {
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
+        server.onError {}
+        server.onMessage {}
+        server.start()
+        server.close()
+
+        val ex = shouldThrow<IllegalStateException> {
+            server.start()
+        }
+        ex.message shouldContain "closed"
+    }
+
     @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 server = StdioServerTransport(input = bufferedInput, output = printOutput)
         val capturedErrors = mutableListOf<Throwable>()
         val receivedMessages = mutableListOf<JSONRPCMessage>()
         val secondMessageProcessed = CompletableDeferred<Unit>()
@@ -253,7 +282,7 @@ class StdioServerTransportTest {
 
     @Test
     fun `should not invoke onError for CancellationException in handler`() = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         val capturedError = CompletableDeferred<Throwable>()
         server.onError { capturedError.complete(it) }
 
@@ -281,7 +310,7 @@ class StdioServerTransportTest {
 
     @Test
     fun `should continue receiving valid messages after malformed JSON is skipped`() = runIntegrationTest {
-        val server = StdioServerTransport(bufferedInput, printOutput)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
         val received = CompletableDeferred<JSONRPCMessage>()
         // ReadBuffer silently skips unparseable lines — no onError callback expected
         server.onError {}
@@ -299,6 +328,297 @@ class StdioServerTransportTest {
         server.close()
     }
 
+    @Test
+    fun `should throw McpException when send is called before start`() = runIntegrationTest {
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
+        val ex = shouldThrow<McpException> {
+            server.send(PingRequest().toJSON())
+        }
+        ex.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED
+    }
+
+    @Test
+    fun `should throw McpException when send is called after close`() = runIntegrationTest {
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
+        server.onError {}
+        server.onMessage {}
+        server.start()
+        server.close()
+
+        val ex = shouldThrow<McpException> {
+            server.send(PingRequest().toJSON())
+        }
+        ex.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED
+    }
+
+    @Test
+    fun `should drain in-flight outbound messages on graceful close`() = runIntegrationTest {
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
+        server.onError {}
+        server.onMessage {}
+        server.start()
+
+        val numMessages = 50
+        repeat(numMessages) {
+            server.send(PingRequest().toJSON())
+        }
+        server.close()
+
+        val outputLines = String(output.toByteArray()).lines().count { it.isNotBlank() }
+        outputLines shouldBe numMessages
+    }
+
+    @Test
+    fun `should fail-fast when reader throws non-EOF IOException`() = runIntegrationTest {
+        val ioError = IOException("transient stream failure")
+        val server = StdioServerTransport(input = FaultyRawSource(ioError).buffered(), output = printOutput)
+        val errorCaptured = CompletableDeferred<Throwable>()
+        val closeCaptured = CompletableDeferred<Unit>()
+        server.onError { errorCaptured.complete(it) }
+        server.onClose { closeCaptured.complete(Unit) }
+        server.onMessage {}
+
+        server.start()
+
+        errorCaptured.await() shouldBe ioError
+        closeCaptured.await()
+    }
+
+    @Test
+    fun `should close input source on natural EOF`() = runIntegrationTest {
+        val inputClosed = CompletableDeferred<Unit>()
+        val eofSource = object : RawSource {
+            override fun readAtMostTo(sink: Buffer, byteCount: Long): Long = -1L
+            override fun close() {
+                inputClosed.complete(Unit)
+            }
+        }
+        val server = StdioServerTransport(input = eofSource.buffered(), output = printOutput)
+        val onCloseFired = CompletableDeferred<Unit>()
+        server.onError {}
+        server.onMessage {}
+        server.onClose { onCloseFired.complete(Unit) }
+
+        server.start()
+
+        inputClosed.await()
+        onCloseFired.await()
+    }
+
+    // region: concurrency
+
+    @Test
+    fun `should be safe under concurrent start and close`() = runIntegrationTest(timeout = 30.seconds) {
+        val iterations = 2000
+        repeat(iterations) { iteration ->
+            val server = StdioServerTransport(input = Buffer(), output = Buffer())
+            server.onError {}
+            server.onMessage {}
+
+            coroutineScope {
+                launch(Dispatchers.Default) {
+                    try {
+                        server.start()
+                    } catch (_: IllegalStateException) {
+                    }
+                }
+                launch(Dispatchers.Default) {
+                    server.close()
+                }
+            }
+
+            withClue("iteration $iteration: send() after race should fail with CONNECTION_CLOSED") {
+                val ex = shouldThrow<McpException> {
+                    server.send(PingRequest().toJSON())
+                }
+                ex.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED
+            }
+        }
+    }
+
+    @Test
+    fun `should suspend send under back-pressure`() = runIntegrationTest(timeout = 30.seconds) {
+        val unblock = CountDownLatch(1)
+        val blockingSink = object : RawSink {
+            override fun write(source: Buffer, byteCount: Long) {
+                source.clear()
+                unblock.await()
+            }
+
+            override fun flush() {
+                // noop
+            }
+
+            override fun close() {
+                unblock.countDown()
+            }
+        }
+        val server = StdioServerTransport(input = bufferedInput, output = blockingSink.buffered())
+        server.onError {}
+        server.onMessage {}
+        server.start()
+
+        val maxAttempts = 256
+        var suspendedAt = -1
+        for (i in 0 until maxAttempts) {
+            val ok = withTimeoutOrNull(100.milliseconds) { server.send(PingRequest().toJSON()) }
+            if (ok == null) {
+                suspendedAt = i
+                break
+            }
+        }
+        assertTrue(
+            suspendedAt in 0 until maxAttempts,
+            "send() should have suspended within $maxAttempts attempts, got suspendedAt=$suspendedAt",
+        )
+        // Unblock writer so close() can drain.
+        unblock.countDown()
+        server.close()
+    }
+
+    @Test
+    fun `should not hang close when handler suspends indefinitely`() = runIntegrationTest {
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
+        val handlerEntered = CompletableDeferred<Unit>()
+        val neverComplete = CompletableDeferred<Unit>()
+        server.onError {}
+        server.onMessage {
+            handlerEntered.complete(Unit)
+            neverComplete.await()
+        }
+        server.start()
+
+        inputWriter.write(serializeMessage(PingRequest().toJSON()))
+        inputWriter.flush()
+
+        handlerEntered.await()
+        server.close()
+    }
+
+    @Test
+    fun `should not hang close when onMessage was never registered`() = runIntegrationTest {
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput)
+        server.onError {}
+        server.start()
+
+        inputWriter.write(serializeMessage(PingRequest().toJSON()))
+        inputWriter.flush()
+
+        delay(200.milliseconds)
+
+        server.close()
+    }
+
+    // endregion
+
+    // region: scope and dispatcher knobs
+
+    @Test
+    fun `should honor externally provided CoroutineScope`() = runIntegrationTest {
+        val externalScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
+        val server = StdioServerTransport(input = bufferedInput, output = printOutput) {
+            scope = externalScope
+        }
+        val closeCaptured = CompletableDeferred<Unit>()
+        server.onError {}
+        server.onClose { closeCaptured.complete(Unit) }
+        server.onMessage {}
+
+        server.start()
+        externalScope.cancel()
+
+        closeCaptured.await()
+
+        val ex = shouldThrow<McpException> {
+            server.send(PingRequest().toJSON())
+        }
+        ex.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED
+    }
+
+    @Test
+    fun `should use ioDispatcher for reader and writer`() = runIntegrationTest {
+        val threadName = "stdio-test-io-thread"
+        val executor = Executors.newFixedThreadPool(2) { r -> Thread(r, threadName) }
+        try {
+            val readerThread = CompletableDeferred<String>()
+            val writerThread = CompletableDeferred<String>()
+            val unblockReader = CountDownLatch(1)
+            val recordingInput = object : RawSource {
+                override fun readAtMostTo(sink: Buffer, byteCount: Long): Long {
+                    if (!readerThread.isCompleted) {
+                        readerThread.complete(Thread.currentThread().name)
+                    }
+                    unblockReader.await()
+                    return -1L
+                }
+
+                override fun close() {
+                    unblockReader.countDown()
+                }
+            }
+            val recordingOutput = object : RawSink {
+                override fun write(source: Buffer, byteCount: Long) {
+                    if (!writerThread.isCompleted) {
+                        writerThread.complete(Thread.currentThread().name)
+                    }
+                    source.clear()
+                }
+
+                override fun flush() {
+                    // noop
+                }
+
+                override fun close() {
+                    // noop
+                }
+            }
+            val server = StdioServerTransport(
+                input = recordingInput.buffered(),
+                output = recordingOutput.buffered(),
+            ) {
+                ioDispatcher = executor.asCoroutineDispatcher()
+            }
+            server.onError {}
+            server.onMessage {}
+            server.start()
+
+            readerThread.await() shouldContain threadName
+            server.send(PingRequest().toJSON())
+            writerThread.await() shouldContain threadName
+
+            server.close()
+        } finally {
+            executor.shutdownNow()
+        }
+    }
+
+    @Test
+    fun `should use handlerDispatcher for message handling`() = runIntegrationTest {
+        val threadName = "stdio-test-handler-thread"
+        val executor = Executors.newSingleThreadExecutor { r -> Thread(r, threadName) }
+        try {
+            val server = StdioServerTransport(input = bufferedInput, output = printOutput) {
+                handlerDispatcher = executor.asCoroutineDispatcher()
+            }
+            val observedThreadName = CompletableDeferred<String>()
+            server.onError {}
+            server.onMessage {
+                observedThreadName.complete(Thread.currentThread().name)
+            }
+            server.start()
+
+            inputWriter.write(serializeMessage(PingRequest().toJSON()))
+            inputWriter.flush()
+
+            observedThreadName.await() shouldContain threadName
+            server.close()
+        } finally {
+            executor.shutdownNow()
+        }
+    }
+
+    // endregion
+
     private fun inputErrors() = listOf(
         IOException("simulated read failure"),
         RuntimeException("unexpected read exception"),