feat: Add kotlin-sdk-testing module with channel-based transport implementation (#505)
# feat: Add kotlin-sdk-testing module with channel-based transport implementation - Introduced `kotlin-sdk-testing` module for testing utilities in the MCP Kotlin SDK. - Implemented `ChannelTransport`, providing JSON-RPC transport via Kotlin coroutines channels. <!-- Provide a brief summary of your changes --> ## Motivation and Context This ChannelTransport can be used for testing Server implementation without using real transports. ## How Has This Been Tested? Unit test ## Breaking Changes No ## Types of changes - [ ] 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) - [x] Documentation update ## Checklist <!-- Go over all the following points, and put an `x` in all the boxes that apply. --> - [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:
Konstantin Pavlov
GitHub
parent
f9d491e4fa
commit
212cbfa325
4 changed files with 669 additions and 3 deletions
|
|
@ -0,0 +1,176 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.testing
|
||||
|
||||
import io.github.oshai.kotlinlogging.KLogger
|
||||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.modelcontextprotocol.kotlin.sdk.ExperimentalMcpApi
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractClientTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import kotlinx.coroutines.CancellationException
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.CoroutineDispatcher
|
||||
import kotlinx.coroutines.CoroutineName
|
||||
import kotlinx.coroutines.CoroutineScope
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.coroutines.Job
|
||||
import kotlinx.coroutines.SupervisorJob
|
||||
import kotlinx.coroutines.cancel
|
||||
import kotlinx.coroutines.channels.Channel
|
||||
import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
|
||||
import kotlinx.coroutines.channels.ReceiveChannel
|
||||
import kotlinx.coroutines.channels.SendChannel
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.yield
|
||||
|
||||
/**
|
||||
* A transport implementation that uses Kotlin Coroutines Channels for asynchronous
|
||||
* communication of JSON-RPC messages. This is useful in scenarios where communication is
|
||||
* message-based and channels are used as the underlying mechanism for delivery.
|
||||
*
|
||||
* @constructor Initializes the transport with the specified [SendChannel] and [ReceiveChannel].
|
||||
* @property sendChannel The channel used to send JSON-RPC messages.
|
||||
* @property receiveChannel The channel used to receive JSON-RPC messages.
|
||||
* @param dispatcher The coroutine dispatcher for the event loop. Defaults to [Dispatchers.Default].
|
||||
*/
|
||||
@ExperimentalMcpApi
|
||||
public class ChannelTransport(
|
||||
private val sendChannel: SendChannel<JSONRPCMessage>,
|
||||
private val receiveChannel: ReceiveChannel<JSONRPCMessage>,
|
||||
dispatcher: CoroutineDispatcher = Dispatchers.Default,
|
||||
) : AbstractClientTransport() {
|
||||
|
||||
override val logger: KLogger = KotlinLogging.logger {}
|
||||
|
||||
private val scope = CoroutineScope(SupervisorJob() + dispatcher)
|
||||
|
||||
/**
|
||||
* Creates a `ChannelTransport` instance using a single channel for both sending and receiving messages.
|
||||
*
|
||||
* This constructor simplifies the use of `ChannelTransport` in scenarios where a single channel
|
||||
* is enough for bidirectional communication. A default channel with unlimited capacity is used
|
||||
* if no channel is provided.
|
||||
*
|
||||
* @param channel The `Channel` of `JSONRPCMessage` instances, used for both sending and receiving.
|
||||
* Defaults to a channel with unlimited capacity.
|
||||
* @param dispatcher The coroutine dispatcher for the event loop. Defaults to [Dispatchers.Default].
|
||||
*/
|
||||
public constructor(
|
||||
channel: Channel<JSONRPCMessage> = Channel(UNLIMITED),
|
||||
dispatcher: CoroutineDispatcher = Dispatchers.Default,
|
||||
) : this(channel, channel, dispatcher)
|
||||
|
||||
/**
|
||||
* Represents a pair of interconnected [ChannelTransport]s for bidirectional communication.
|
||||
*
|
||||
* Messages sent via one transport are received by the other. This is useful for
|
||||
* setting up communication between a client and server without requiring external networking.
|
||||
*
|
||||
* @property clientTransport The transport intended for use on the client-side.
|
||||
* @property serverTransport The transport intended for use on the server-side.
|
||||
*/
|
||||
public data class LinkedTransports(val clientTransport: ChannelTransport, val serverTransport: ChannelTransport)
|
||||
|
||||
public companion object {
|
||||
|
||||
/**
|
||||
* Creates a pair of interconnected [ChannelTransport] objects for bidirectional communication.
|
||||
*
|
||||
* This method sets up a client and server transport using Kotlin channels, enabling messages sent
|
||||
* through one transport to be received by the other. It is typically used for testing or in-memory
|
||||
* communication without external networking.
|
||||
*
|
||||
* @param capacity The buffer capacity of the internal channels used for communication. Defaults to 256.
|
||||
* @return A [LinkedTransports] instance containing the client and server transports.
|
||||
*/
|
||||
public fun createLinkedPair(
|
||||
capacity: Int = 256,
|
||||
dispatcher: CoroutineDispatcher = Dispatchers.Default,
|
||||
): LinkedTransports {
|
||||
val sendChannel = Channel<JSONRPCMessage>(capacity)
|
||||
val receiveChannel = Channel<JSONRPCMessage>(capacity)
|
||||
val clientTransport = ChannelTransport(sendChannel, receiveChannel, dispatcher)
|
||||
val serverTransport = ChannelTransport(receiveChannel, sendChannel, dispatcher)
|
||||
return LinkedTransports(clientTransport = clientTransport, serverTransport = serverTransport)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Starts the transport, allowing it to begin processing messages.
|
||||
*
|
||||
* Waits for the event loop to start before returning, ensuring the transport is
|
||||
* fully initialized and ready to send/receive messages.
|
||||
*
|
||||
* Processes messages until the channel closes or cancellation occurs.
|
||||
* Exceptions from message handlers are reported via onError but do not stop processing.
|
||||
*/
|
||||
override suspend fun initialize() {
|
||||
logger.info { "ChannelTransport starting message processing" }
|
||||
val started = CompletableDeferred<Unit>()
|
||||
@Suppress("TooGenericExceptionCaught")
|
||||
scope.launch(CoroutineName("ChannelTransport#${hashCode()}-event-loop")) {
|
||||
try {
|
||||
// Signal that event loop has started, then yield to ensure we're ready
|
||||
started.complete(Unit)
|
||||
yield()
|
||||
|
||||
for (message in receiveChannel) {
|
||||
logger.debug { "Received message: ${message::class.simpleName}" }
|
||||
|
||||
try {
|
||||
_onMessage.invoke(message)
|
||||
logger.trace { "Message processed successfully: ${message::class.simpleName}" }
|
||||
} catch (e: CancellationException) {
|
||||
// Let cancellation propagate immediately
|
||||
logger.debug { "Cancellation requested during message processing" }
|
||||
throw e
|
||||
} catch (e: Exception) {
|
||||
// Report other errors but continue processing
|
||||
logger.warn(e) { "Error processing message: ${message::class.simpleName}" }
|
||||
_onError.invoke(e)
|
||||
}
|
||||
}
|
||||
logger.info { "ChannelTransport stopped: receive channel closed" }
|
||||
} catch (e: Exception) {
|
||||
// Only complete exceptionally if not already completed
|
||||
if (!started.isCompleted) {
|
||||
started.completeExceptionally(e)
|
||||
}
|
||||
throw e
|
||||
} finally {
|
||||
close()
|
||||
}
|
||||
}
|
||||
// Wait for the event loop to start
|
||||
started.await()
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a JSON-RPC message through the transport.
|
||||
*
|
||||
* @param message The JSON-RPC message to send. This can be a request or a response,
|
||||
* conforming to the JSON-RPC 2.0 specification.
|
||||
* @param options Optional transport-specific options that provide additional context or behavior
|
||||
* for the message transmission, such as associating the message with a specific
|
||||
* incoming request or handling resumption tokens.
|
||||
*/
|
||||
override suspend fun performSend(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
logger.debug { "Sending message: ${message::class.simpleName}" }
|
||||
sendChannel.send(message)
|
||||
logger.trace { "Message sent successfully: ${message::class.simpleName}" }
|
||||
}
|
||||
|
||||
/**
|
||||
* Closes the transport, preventing further message transmission and processing.
|
||||
*/
|
||||
override suspend fun closeResources() {
|
||||
logger.info { "Closing ChannelTransport" }
|
||||
sendChannel.close()
|
||||
if (receiveChannel !== sendChannel) {
|
||||
logger.debug { "Cancelling separate receive channel" }
|
||||
receiveChannel.cancel()
|
||||
}
|
||||
scope.cancel()
|
||||
scope.coroutineContext[Job]?.join() // Wait for cleanup
|
||||
logger.info { "ChannelTransport closed" }
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,212 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.testing
|
||||
|
||||
import io.kotest.assertions.nondeterministic.eventually
|
||||
import io.kotest.matchers.collections.shouldContainExactly
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.kotest.matchers.types.shouldBeInstanceOf
|
||||
import io.modelcontextprotocol.kotlin.sdk.ExperimentalMcpApi
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.cancelAndJoin
|
||||
import kotlinx.coroutines.channels.Channel
|
||||
import kotlinx.coroutines.channels.ClosedSendChannelException
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlin.test.Test
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
@OptIn(ExperimentalMcpApi::class)
|
||||
class ChannelTransportTest {
|
||||
|
||||
@Test
|
||||
fun `send writes to sendChannel`() = runTest {
|
||||
val sendChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(sendChannel, Channel())
|
||||
|
||||
transport.start()
|
||||
|
||||
val message = JSONRPCRequest(RequestId.NumberId(42), "test")
|
||||
transport.send(message)
|
||||
|
||||
sendChannel.receive() shouldBe message
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `start reads from receiveChannel and invokes onMessage`() = runTest {
|
||||
val receiveChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(Channel(), receiveChannel)
|
||||
|
||||
val messagesProcessed = CompletableDeferred<Unit>()
|
||||
val received = mutableListOf<JSONRPCMessage>()
|
||||
transport.onMessage { msg ->
|
||||
received.add(msg)
|
||||
if (received.size == 2) {
|
||||
messagesProcessed.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
transport.start()
|
||||
|
||||
val msg1 = JSONRPCRequest(RequestId.NumberId(1), "method1")
|
||||
val msg2 = JSONRPCRequest(RequestId.NumberId(2), "method2")
|
||||
receiveChannel.send(msg1)
|
||||
receiveChannel.send(msg2)
|
||||
|
||||
// Wait for messages to be processed
|
||||
messagesProcessed.await()
|
||||
|
||||
received.shouldContainExactly(msg1, msg2)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `start completes when channel closes`() = runTest {
|
||||
val receiveChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(Channel(), receiveChannel)
|
||||
|
||||
val messageProcessed = CompletableDeferred<Unit>()
|
||||
val received = mutableListOf<JSONRPCMessage>()
|
||||
transport.onMessage {
|
||||
received.add(it)
|
||||
messageProcessed.complete(Unit)
|
||||
}
|
||||
|
||||
transport.start()
|
||||
|
||||
val msg = JSONRPCRequest(RequestId.NumberId(1), "method1")
|
||||
receiveChannel.send(msg)
|
||||
receiveChannel.close() // Close the channel
|
||||
|
||||
// Wait for a message to be processed
|
||||
messageProcessed.await()
|
||||
|
||||
received.shouldContainExactly(msg)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `close is idempotent`() = runTest {
|
||||
val transport = ChannelTransport()
|
||||
|
||||
var closeCount = 0
|
||||
transport.onClose { closeCount++ }
|
||||
|
||||
transport.start()
|
||||
|
||||
transport.close()
|
||||
transport.close()
|
||||
|
||||
closeCount shouldBe 1
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `callbacks chain not replace`() = runTest {
|
||||
val receiveChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(Channel(), receiveChannel)
|
||||
|
||||
val messageProcessed = CompletableDeferred<Unit>()
|
||||
val calls = mutableListOf<String>()
|
||||
transport.onMessage {
|
||||
calls.add("first")
|
||||
}
|
||||
transport.onMessage {
|
||||
calls.add("second")
|
||||
messageProcessed.complete(Unit)
|
||||
}
|
||||
|
||||
val startJob = backgroundScope.launch { transport.start() }
|
||||
|
||||
receiveChannel.send(JSONRPCRequest(RequestId.NumberId(1), "test"))
|
||||
messageProcessed.await()
|
||||
|
||||
calls.shouldContainExactly("first", "second")
|
||||
startJob.cancelAndJoin()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `secondary constructor uses single channel`() = runTest {
|
||||
val channel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(channel)
|
||||
|
||||
val messageProcessed = CompletableDeferred<Unit>()
|
||||
val received = mutableListOf<JSONRPCMessage>()
|
||||
transport.onMessage { msg ->
|
||||
received.add(msg)
|
||||
messageProcessed.complete(Unit)
|
||||
}
|
||||
|
||||
transport.start()
|
||||
|
||||
val message = JSONRPCRequest(RequestId.NumberId(1), "test")
|
||||
transport.send(message)
|
||||
messageProcessed.await()
|
||||
|
||||
eventually(2.seconds) {
|
||||
received.shouldContainExactly(message)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `send to closed channel triggers error`() = runTest {
|
||||
val sendChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(sendChannel, Channel())
|
||||
|
||||
transport.start()
|
||||
|
||||
var errorCaught = false
|
||||
transport.onError {
|
||||
it.shouldBeInstanceOf<ClosedSendChannelException>()
|
||||
errorCaught = true
|
||||
}
|
||||
sendChannel.close()
|
||||
|
||||
try {
|
||||
transport.send(JSONRPCRequest(RequestId.NumberId(1), "method"))
|
||||
} catch (e: Exception) {
|
||||
// send() wraps ClosedSendChannelException in McpException
|
||||
e.shouldBeInstanceOf<io.modelcontextprotocol.kotlin.sdk.types.McpException>()
|
||||
}
|
||||
|
||||
eventually(2.seconds) {
|
||||
errorCaught shouldBe true
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `exceptions in message handler do not stop processing`() = runTest {
|
||||
val receiveChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
val transport = ChannelTransport(Channel(), receiveChannel)
|
||||
|
||||
val allProcessed = CompletableDeferred<Unit>()
|
||||
val received = mutableListOf<Int>()
|
||||
val errors = mutableListOf<Throwable>()
|
||||
|
||||
transport.onError { errors.add(it) }
|
||||
transport.onMessage { msg ->
|
||||
val id = ((msg as JSONRPCRequest).id as RequestId.NumberId).value.toInt()
|
||||
received.add(id)
|
||||
if (id == 2) {
|
||||
@Suppress("TooGenericExceptionThrown")
|
||||
throw RuntimeException("Error processing message 2")
|
||||
}
|
||||
if (received.size == 4) {
|
||||
allProcessed.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
transport.start()
|
||||
|
||||
// Send 4 messages, second one will throw
|
||||
receiveChannel.send(JSONRPCRequest(RequestId.NumberId(1), "m1"))
|
||||
receiveChannel.send(JSONRPCRequest(RequestId.NumberId(2), "m2"))
|
||||
receiveChannel.send(JSONRPCRequest(RequestId.NumberId(3), "m3"))
|
||||
receiveChannel.send(JSONRPCRequest(RequestId.NumberId(4), "m4"))
|
||||
|
||||
allProcessed.await()
|
||||
|
||||
// All messages should be processed despite error in message 2
|
||||
received.shouldContainExactly(1, 2, 3, 4)
|
||||
errors.size shouldBe 1
|
||||
errors[0].message shouldBe "Error processing message 2"
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
## https://docs.junit.org/5.3.0-M1/user-guide/index.html#writing-tests-parallel-execution
|
||||
junit.jupiter.execution.parallel.enabled=true
|
||||
junit.jupiter.execution.parallel.config.strategy=dynamic
|
||||
junit.jupiter.execution.parallel.mode.default=concurrent
|
||||
junit.jupiter.execution.parallel.mode.classes.default=concurrent
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
# Level of logging for the ROOT logger: ERROR, WARN, INFO, DEBUG, TRACE (default is INFO)
|
||||
org.slf4j.simpleLogger.defaultLogLevel=INFO
|
||||
org.slf4j.simpleLogger.showThreadName=true
|
||||
org.slf4j.simpleLogger.showDateTime=false
|
||||
|
||||
# Log level for specific packages or classes
|
||||
org.slf4j.simpleLogger.log.io.modelcontextprotocol=DEBUG
|
||||
Loading…
Reference in a new issue