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
This commit is contained in:
Pavel Gorgulov
GitHub
parent
37ee423308
commit
0adfb3bcdb
4 changed files with 642 additions and 147 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -56,7 +56,7 @@ node_modules
|
|||
dist
|
||||
|
||||
### SWE agents ###
|
||||
.claude/settings.local.json
|
||||
.claude/
|
||||
.junie/
|
||||
|
||||
### Conformance test results ###
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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].
|
||||
* @param inputStream The input [Source] used to receive data.
|
||||
* @param outputStream The output [Sink] used to send data.
|
||||
* - **reader** — pulls bytes from the input source into the parsing buffer; runs on
|
||||
* [Builder.ioDispatcher].
|
||||
* - **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 var readingJob: Job? = null
|
||||
private var sendingJob: Job? = null
|
||||
private var processingJob: Job? = null
|
||||
private val readChannel = Channel<ByteArray>(Channel.BUFFERED)
|
||||
private val writeChannel = Channel<JSONRPCMessage>(Channel.BUFFERED)
|
||||
private val outputSink = output.buffered()
|
||||
|
||||
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()
|
||||
private enum class State { New, Operational, Stopped }
|
||||
|
||||
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)) {
|
||||
error("StdioServerTransport already started!")
|
||||
if (!state.compareAndSet(State.New, State.Operational)) {
|
||||
when (state.load()) {
|
||||
State.Stopped -> error("StdioServerTransport is already closed!")
|
||||
else -> error("StdioServerTransport already started!")
|
||||
}
|
||||
}
|
||||
|
||||
// Launch a coroutine to read from stdin
|
||||
readingJob = launchReadingJob()
|
||||
try {
|
||||
val resolvedScope = scope ?: CoroutineScope(
|
||||
currentCoroutineContext() + IODispatcher + SupervisorJob(),
|
||||
)
|
||||
effectiveScope = resolvedScope
|
||||
|
||||
// Launch a coroutine to process messages from readChannel
|
||||
processingJob = launchProcessingJob()
|
||||
|
||||
// Launch a coroutine to handle message sending
|
||||
sendingJob = launchSendingJob()
|
||||
readerJob = resolvedScope.launch(
|
||||
ioDispatcher + CoroutineName("StdioServerTransport.reader"),
|
||||
) { readerPump() }
|
||||
processorJob = resolvedScope.launch(
|
||||
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 {
|
||||
val job = scope.launch {
|
||||
val buf = Buffer()
|
||||
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 suspend fun readerPump() {
|
||||
val buf = Buffer()
|
||||
try {
|
||||
while (currentCoroutineContext().isActive) {
|
||||
val bytesRead = input.readAtMostTo(buf, READ_BUFFER_SIZE)
|
||||
if (bytesRead == -1L) break
|
||||
if (bytesRead > 0) {
|
||||
val chunk = buf.readByteArray()
|
||||
readChannel.send(chunk)
|
||||
}
|
||||
}
|
||||
} catch (e: CancellationException) {
|
||||
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?) {
|
||||
when (cause) {
|
||||
is CancellationException -> {
|
||||
}
|
||||
|
||||
null -> {
|
||||
logger.debug { "$jobName job completed" }
|
||||
}
|
||||
|
||||
else -> {
|
||||
logger.debug(cause) { "$jobName job completed exceptionally" }
|
||||
private suspend fun writerPump() {
|
||||
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 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()
|
||||
sendingJob?.cancelAndJoin()
|
||||
setupComplete.await()
|
||||
|
||||
runCatching {
|
||||
inputStream.close()
|
||||
}.onFailure { logger.warn(it) { "Failed to close stdin" } }
|
||||
if (previous == State.Stopped) {
|
||||
writerJob?.join()
|
||||
return@withContext
|
||||
}
|
||||
|
||||
readingJob?.cancel()
|
||||
runCatching { input.close() }
|
||||
.onFailure { logger.warn(it) { "Failed to close input source" } }
|
||||
readerJob?.cancel()
|
||||
readChannel.close()
|
||||
|
||||
processingJob?.cancelAndJoin()
|
||||
|
||||
processorJob?.cancelAndJoin()
|
||||
writeChannel.close()
|
||||
writerJob?.join()
|
||||
runCatching { outputSink.close() }
|
||||
.onFailure { logger.warn(it) { "Failed to close output sink" } }
|
||||
readBuffer.clear()
|
||||
|
||||
runCatching {
|
||||
outputSink.flush()
|
||||
outputSink.close()
|
||||
}.onFailure { logger.warn(it) { "Failed to close stdout" } }
|
||||
|
||||
if (ownsScope) {
|
||||
effectiveScope?.coroutineContext?.get(Job)?.cancel()
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
Loading…
Reference in a new issue