modelcontextprotocol/kotlin-sdk
Watch
1
0
Fork
You've already forked kotlin-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Modify SSE config for server (#23)

Browse source 
* Fixed compilation

* Added `mcp` Ktor server route alongside Application extension

* Fix AbstractTransport to not skip messages
This commit is contained in:
Alexander Sysoev 2025-01-24 13:31:32 +01:00 • committed by GitHub
commit de5ef4460f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 296 additions and 355 deletions
Download patch file Download diff file Expand all files Collapse all files

34
README.md
View file

@ -114,12 +114,13 @@ server.connect(transport)
### Using SSE Transport
Directly in Ktor's `Application`:
```kotlin
import io.ktor.server.application.*
import io.modelcontextprotocol.kotlin.sdk.server.MCP
import io.modelcontextprotocol.kotlin.sdk.server.mcp
fun Application.module() {
MCP {
mcp {
Server(
serverInfo = Implementation(
name = "example-sse-server",
@ -136,6 +137,35 @@ fun Application.module() {
}
```
Inside a custom Ktor's `Route`:
```kotlin
import io.ktor.server.application.*
import io.ktor.server.sse.SSE
import io.modelcontextprotocol.kotlin.sdk.server.mcp
fun Application.module() {
install(SSE)
routing {
route("myRoute") {
mcp {
Server(
serverInfo = Implementation(
name = "example-sse-server",
version = "1.0.0"
),
options = ServerOptions(
capabilities = ServerCapabilities(
prompts = ServerCapabilities.Prompts(listChanged = null),
resources = ServerCapabilities.Resources(subscribe = null, listChanged = null)
)
)
)
}
}
}
}
```
## Contributing
Please see the [contribution guide](CONTRIBUTING.md) and the [Code of conduct](CODE_OF_CONDUCT.md) before contributing.

1
build.gradle.kts
View file

@ -225,6 +225,7 @@ kotlin {
jvmTest {
dependencies {
implementation(libs.mockk)
implementation(libs.slf4j.simple)
}
}
}

2
gradle/libs.versions.toml
View file

@ -12,6 +12,7 @@ mockk = "1.13.13"
logging = "7.0.0"
jreleaser = "1.15.0"
binaryCompatibilityValidatorPlugin = "0.17.0"
slf4j = "2.0.16"
[libraries]
# Kotlinx libraries
@ -30,6 +31,7 @@ kotlinx-coroutines-test = { group = "org.jetbrains.kotlinx", name = "kotlinx-cor
kotlinx-coroutines-debug = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-debug", version.ref = "coroutines" }
ktor-server-test-host = { group = "io.ktor", name = "ktor-server-test-host", version.ref = "ktor" }
mockk = { group = "io.mockk", name = "mockk", version.ref = "mockk" }
slf4j-simple = { group = "org.slf4j", name = "slf4j-simple", version.ref = "slf4j" }

46
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/KtorClient.kt Normal file
View file

@ -0,0 +1,46 @@
package io.modelcontextprotocol.kotlin.sdk.client
import io.ktor.client.HttpClient
import io.ktor.client.request.HttpRequestBuilder
import io.modelcontextprotocol.kotlin.sdk.Implementation
import io.modelcontextprotocol.kotlin.sdk.LIB_VERSION
import io.modelcontextprotocol.kotlin.sdk.shared.IMPLEMENTATION_NAME
import kotlin.time.Duration
/**
* Returns a new SSE transport for the Model Context Protocol using the provided HttpClient.
*
* @param urlString Optional URL of the MCP server.
* @param reconnectionTime Optional duration to wait before attempting to reconnect.
* @param requestBuilder Optional lambda to configure the HTTP request.
* @return A [SSEClientTransport] configured for MCP communication.
*/
public fun HttpClient.mcpSseTransport(
urlString: String? = null,
reconnectionTime: Duration? = null,
requestBuilder: HttpRequestBuilder.() -> Unit = {},
): SseClientTransport = SseClientTransport(this, urlString, reconnectionTime, requestBuilder)
/**
* Creates and connects an MCP client over SSE using the provided HttpClient.
*
* @param urlString Optional URL of the MCP server.
* @param reconnectionTime Optional duration to wait before attempting to reconnect.
* @param requestBuilder Optional lambda to configure the HTTP request.
* @return A connected [Client] ready for MCP communication.
*/
public suspend fun HttpClient.mcpSse(
urlString: String? = null,
reconnectionTime: Duration? = null,
requestBuilder: HttpRequestBuilder.() -> Unit = {},
): Client {
val transport = mcpSseTransport(urlString, reconnectionTime, requestBuilder)
val client = Client(
Implementation(
name = IMPLEMENTATION_NAME,
version = LIB_VERSION
)
)
client.connect(transport)
return client
}

37
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt
View file

@ -6,8 +6,8 @@ import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.http.*
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import io.modelcontextprotocol.kotlin.sdk.shared.McpJson
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
import kotlinx.atomicfu.AtomicBoolean
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.*
@ -15,16 +15,19 @@ import kotlinx.serialization.encodeToString
import kotlin.properties.Delegates
import kotlin.time.Duration
@Deprecated("Use SseClientTransport instead", ReplaceWith("SseClientTransport"), DeprecationLevel.WARNING)
public typealias SSEClientTransport = SseClientTransport
/**
* Client transport for SSE: this will connect to a server using Server-Sent Events for receiving
* messages and make separate POST requests for sending messages.
*/
public class SSEClientTransport(
public class SseClientTransport(
private val client: HttpClient,
private val urlString: String?,
private val reconnectionTime: Duration? = null,
private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
) : Transport {
) : AbstractTransport() {
private val scope by lazy {
CoroutineScope(session.coroutineContext + SupervisorJob())
}
@ -33,10 +36,6 @@ public class SSEClientTransport(
private var session: ClientSSESession by Delegates.notNull()
private val endpoint = CompletableDeferred<String>()
private var _onClose: (() -> Unit) = {}
private var _onError: ((Throwable) -> Unit) = {}
private var _onMessage: (suspend ((JSONRPCMessage) -> Unit)) = {}
private var job: Job? = null
private val baseUrl by lazy {
@ -136,28 +135,4 @@ public class SSEClientTransport(
_onClose()
job?.cancelAndJoin()
}
override fun onClose(block: () -> Unit) {
val old = _onClose
_onClose = {
old()
block()
}
}
override fun onError(block: (Throwable) -> Unit) {
val old = _onError
_onError = { e ->
old(e)
block(e)
}
}
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {
val old = _onMessage
_onMessage = { message ->
old(message)
block(message)
}
}
}

20
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport.kt
View file

@ -2,8 +2,8 @@ package io.modelcontextprotocol.kotlin.sdk.client
import io.github.oshai.kotlinlogging.KotlinLogging
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
import kotlinx.atomicfu.AtomicBoolean
import kotlinx.atomicfu.atomic
@ -30,7 +30,7 @@ import kotlin.coroutines.CoroutineContext
public class StdioClientTransport(
private val input: Source,
private val output: Sink
) : Transport {
) : AbstractTransport() {
private val logger = KotlinLogging.logger {}
private val ioCoroutineContext: CoroutineContext = Dispatchers.IO
private val scope by lazy {
@ -41,10 +41,6 @@ public class StdioClientTransport(
private val sendChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
private val readBuffer = ReadBuffer()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
override suspend fun start() {
if (!initialized.compareAndSet(false, true)) {
error("StdioClientTransport already started!")
@ -70,7 +66,7 @@ public class StdioClientTransport(
}
}
} catch (e: Exception) {
onError?.invoke(e)
_onError.invoke(e)
logger.error(e) { "Error reading from input stream" }
}
}
@ -85,7 +81,7 @@ public class StdioClientTransport(
}
} catch (e: Throwable) {
if (isActive) {
onError?.invoke(e)
_onError.invoke(e)
logger.error(e) { "Error writing to output stream" }
}
} finally {
@ -95,7 +91,7 @@ public class StdioClientTransport(
readJob.join()
writeJob.cancelAndJoin()
onClose?.invoke()
_onClose.invoke()
}
}
@ -116,16 +112,16 @@ public class StdioClientTransport(
output.close()
readBuffer.clear()
sendChannel.close()
onClose?.invoke()
_onClose.invoke()
}
private suspend fun processReadBuffer() {
while (true) {
val msg = readBuffer.readMessage() ?: break
try {
onMessage?.invoke(msg)
_onMessage.invoke(msg)
} catch (e: Throwable) {
onError?.invoke(e)
_onError.invoke(e)
logger.error(e) { "Error processing message." }
}
}

110
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorServer.kt Normal file
View file

@ -0,0 +1,110 @@
package io.modelcontextprotocol.kotlin.sdk.server
import io.github.oshai.kotlinlogging.KotlinLogging
import io.ktor.http.*
import io.ktor.server.application.*
import io.ktor.server.response.*
import io.ktor.server.routing.*
import io.ktor.server.sse.*
import io.ktor.util.collections.*
import io.ktor.utils.io.KtorDsl
private val logger = KotlinLogging.logger {}
@KtorDsl
public fun Routing.mcp(path: String, block: () -> Server) {
route(path) {
mcp(block)
}
}
/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over Server-Sent Events (SSE).
*/
@KtorDsl
public fun Routing.mcp(block: () -> Server) {
val transports = ConcurrentMap<String, SseServerTransport>()
sse {
mcpSseEndpoint("", transports, block)
}
post {
mcpPostEndpoint(transports)
}
}
@Suppress("FunctionName")
@Deprecated("Use mcp() instead", ReplaceWith("mcp(block)"), DeprecationLevel.WARNING)
public fun Application.MCP(block: () -> Server) {
mcp(block)
}
@KtorDsl
public fun Application.mcp(block: () -> Server) {
val transports = ConcurrentMap<String, SseServerTransport>()
install(SSE)
routing {
sse("/sse") {
mcpSseEndpoint("/message", transports, block)
}
post("/message") {
mcpPostEndpoint(transports)
}
}
}
private suspend fun ServerSSESession.mcpSseEndpoint(
postEndpoint: String,
transports: ConcurrentMap<String, SseServerTransport>,
block: () -> Server,
) {
val transport = mcpSseTransport(postEndpoint, transports)
val server = block()
server.onClose {
logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
transports.remove(transport.sessionId)
}
server.connect(transport)
logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" }
}
internal fun ServerSSESession.mcpSseTransport(
postEndpoint: String,
transports: ConcurrentMap<String, SseServerTransport>,
): SseServerTransport {
val transport = SseServerTransport(postEndpoint, this)
transports[transport.sessionId] = transport
logger.info { "New SSE connection established and stored with sessionId: ${transport.sessionId}" }
return transport
}
internal suspend fun RoutingContext.mcpPostEndpoint(
transports: ConcurrentMap<String, SseServerTransport>,
) {
val sessionId: String = call.request.queryParameters["sessionId"]
?: run {
call.respond(HttpStatusCode.BadRequest, "sessionId query parameter is not provided")
return
}
logger.debug { "Received message for sessionId: $sessionId" }
val transport = transports[sessionId]
if (transport == null) {
logger.warn { "Session not found for sessionId: $sessionId" }
call.respond(HttpStatusCode.NotFound, "Session not found")
return
}
transport.handlePostMessage(call)
logger.trace { "Message handled for sessionId: $sessionId" }
}

54
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/McpKtorServerPlugin.kt
View file

@ -1,54 +0,0 @@
package io.modelcontextprotocol.kotlin.sdk.server
import io.github.oshai.kotlinlogging.KotlinLogging
import io.ktor.http.*
import io.ktor.server.application.*
import io.ktor.server.response.*
import io.ktor.server.routing.*
import io.ktor.server.sse.*
import io.ktor.util.collections.*
private val logger = KotlinLogging.logger {}
/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over Server-Sent Events (SSE).
*/
public fun Application.MCP(block: () -> Server) {
val servers = ConcurrentMap<String, Server>()
install(SSE)
routing {
sse("/sse") {
val transport = SSEServerTransport("/message", this)
logger.info { "New SSE connection established with sessionId: ${transport.sessionId}" }
val server = block()
servers[transport.sessionId] = server
logger.debug { "Server instance created and stored for sessionId: ${transport.sessionId}" }
server.onCloseCallback = {
logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
servers.remove(transport.sessionId)
}
server.connect(transport)
logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" }
}
post("/message") {
val sessionId: String = call.request.queryParameters["sessionId"]!!
logger.debug { "Received message for sessionId: $sessionId" }
val transport = servers[sessionId]?.transport as? SSEServerTransport
if (transport == null) {
logger.warn { "Session not found for sessionId: $sessionId" }
call.respond(HttpStatusCode.NotFound, "Session not found")
return@post
}
transport.handlePostMessage(call)
logger.trace { "Message handled for sessionId: $sessionId" }
}
}
}

25
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/SSEServerTransport.kt
View file

@ -6,8 +6,8 @@ import io.ktor.server.request.*
import io.ktor.server.response.*
import io.ktor.server.sse.*
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import io.modelcontextprotocol.kotlin.sdk.shared.McpJson
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
import kotlinx.atomicfu.AtomicBoolean
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.job
@ -17,24 +17,23 @@ import kotlin.uuid.Uuid
internal const val SESSION_ID_PARAM = "sessionId"
@Deprecated("Use SseServerTransport instead", ReplaceWith("SseServerTransport"), DeprecationLevel.WARNING)
public typealias SSEServerTransport = SseServerTransport
/**
* Server transport for SSE: this will send messages over an SSE connection and receive messages from HTTP POST requests.
*
* Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL identified by `_endpoint`.
*/
public class SSEServerTransport(
public class SseServerTransport(
private val endpoint: String,
private val session: ServerSSESession,
) : Transport {
) : AbstractTransport() {
private val initialized: AtomicBoolean = atomic(false)
@OptIn(ExperimentalUuidApi::class)
public val sessionId: String = Uuid.random().toString()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
/**
* Handles the initial SSE connection request.
*
@ -54,7 +53,7 @@ public class SSEServerTransport(
try {
session.coroutineContext.job.join()
} finally {
onClose?.invoke()
_onClose.invoke()
}
}
@ -67,7 +66,7 @@ public class SSEServerTransport(
if (!initialized.value) {
val message = "SSE connection not established"
call.respondText(message, status = HttpStatusCode.InternalServerError)
onError?.invoke(IllegalStateException(message))
_onError.invoke(IllegalStateException(message))
}
val body = try {
@ -79,7 +78,7 @@ public class SSEServerTransport(
call.receiveText()
} catch (e: Exception) {
call.respondText("Invalid message: ${e.message}", status = HttpStatusCode.BadRequest)
onError?.invoke(e)
_onError.invoke(e)
return
}
@ -100,16 +99,16 @@ public class SSEServerTransport(
public suspend fun handleMessage(message: String) {
try {
val parsedMessage = McpJson.decodeFromString<JSONRPCMessage>(message)
onMessage?.invoke(parsedMessage)
_onMessage.invoke(parsedMessage)
} catch (e: Exception) {
onError?.invoke(e)
_onError.invoke(e)
throw e
}
}
override suspend fun close() {
session.close()
onClose?.invoke()
_onClose.invoke()
}
override suspend fun send(message: JSONRPCMessage) {

19
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt
View file

@ -2,8 +2,8 @@ package io.modelcontextprotocol.kotlin.sdk.server
import io.github.oshai.kotlinlogging.KotlinLogging
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
import kotlinx.atomicfu.AtomicBoolean
import kotlinx.atomicfu.atomic
@ -27,11 +27,8 @@ import kotlin.coroutines.CoroutineContext
public class StdioServerTransport(
private val inputStream: Source, //BufferedInputStream = BufferedInputStream(System.`in`),
outputStream: Sink //PrintStream = System.out
) : Transport {
) : AbstractTransport() {
private val logger = KotlinLogging.logger {}
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null
private val readBuffer = ReadBuffer()
private val initialized: AtomicBoolean = atomic(false)
@ -65,7 +62,7 @@ public class StdioServerTransport(
}
} catch (e: Throwable) {
logger.error(e) { "Error reading from stdin" }
onError?.invoke(e)
_onError.invoke(e)
} finally {
// Reached EOF or error, close connection
close()
@ -80,7 +77,7 @@ public class StdioServerTransport(
processReadBuffer()
}
} catch (e: Throwable) {
onError?.invoke(e)
_onError.invoke(e)
}
}
}
@ -90,16 +87,16 @@ public class StdioServerTransport(
val message = try {
readBuffer.readMessage()
} catch (e: Throwable) {
onError?.invoke(e)
_onError.invoke(e)
null
}
if (message == null) break
// Async invocation broke delivery order
try {
onMessage?.invoke(message)
_onMessage.invoke(message)
} catch (e: Throwable) {
onError?.invoke(e)
_onError.invoke(e)
}
}
}
@ -112,7 +109,7 @@ public class StdioServerTransport(
readChannel.close()
readBuffer.clear()
onClose?.invoke()
_onClose.invoke()
}
override suspend fun send(message: JSONRPCMessage) {

8
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt
View file

@ -152,15 +152,15 @@ public abstract class Protocol(
*/
public open suspend fun connect(transport: Transport) {
this.transport = transport
transport.onClose = {
transport.onClose {
doClose()
}
transport.onError = {
transport.onError {
onError(it)
}
transport.onMessage = { message ->
transport.onMessage { message ->
when (message) {
is JSONRPCResponse -> onResponse(message, null)
is JSONRPCRequest -> onRequest(message)
@ -477,4 +477,4 @@ public abstract class Protocol(
public fun removeNotificationHandler(method: Method) {
notificationHandlers.remove(method.value)
}
}
}

51
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt
View file

@ -1,6 +1,7 @@
package io.modelcontextprotocol.kotlin.sdk.shared
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import kotlinx.coroutines.CompletableDeferred
/**
* Describes the minimal contract for a MCP transport that a client or server can communicate over.
@ -46,3 +47,53 @@ public interface Transport {
*/
public fun onMessage(block: suspend (JSONRPCMessage) -> Unit)
}
/**
* Implements [onClose], [onError] and [onMessage] functions of [Transport] providing
* corresponding [_onClose], [_onError] and [_onMessage] properties to use for an implementation.
*/
@Suppress("PropertyName")
public abstract class AbstractTransport : Transport {
protected var _onClose: (() -> Unit) = {}
private set
protected var _onError: ((Throwable) -> Unit) = {}
private set
// to not skip messages
private val _onMessageInitialized = CompletableDeferred<Unit>()
protected var _onMessage: (suspend ((JSONRPCMessage) -> Unit)) = {
_onMessageInitialized.await()
_onMessage.invoke(it)
}
private set
override fun onClose(block: () -> Unit) {
val old = _onClose
_onClose = {
old()
block()
}
}
override fun onError(block: (Throwable) -> Unit) {
val old = _onError
_onError = { e ->
old(e)
block(e)
}
}
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {
val old: suspend (JSONRPCMessage) -> Unit = when (_onMessageInitialized.isCompleted) {
true -> _onMessage
false -> { _ -> }
}
_onMessage = { message ->
old(message)
block(message)
}
_onMessageInitialized.complete(Unit)
}
}

16
src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt
View file

@ -22,7 +22,7 @@ internal const val MCP_SUBPROTOCOL = "mcp"
* Abstract class representing a WebSocket transport for the Model Context Protocol (MCP).
* Handles communication over a WebSocket session.
*/
public abstract class WebSocketMcpTransport : Transport {
public abstract class WebSocketMcpTransport : AbstractTransport() {
private val scope by lazy {
CoroutineScope(session.coroutineContext + SupervisorJob())
}
@ -33,10 +33,6 @@ public abstract class WebSocketMcpTransport : Transport {
*/
protected abstract val session: WebSocketSession
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
/**
* Initializes the WebSocket session
*/
@ -62,15 +58,15 @@ public abstract class WebSocketMcpTransport : Transport {
if (message !is Frame.Text) {
val e = IllegalArgumentException("Expected text frame, got ${message::class.simpleName}: $message")
onError?.invoke(e)
_onError.invoke(e)
throw e
}
try {
val message = McpJson.decodeFromString<JSONRPCMessage>(message.readText())
onMessage?.invoke(message)
_onMessage.invoke(message)
} catch (e: Exception) {
onError?.invoke(e)
_onError.invoke(e)
throw e
}
}
@ -79,9 +75,9 @@ public abstract class WebSocketMcpTransport : Transport {
@OptIn(InternalCoroutinesApi::class)
session.coroutineContext.job.invokeOnCompletion {
if (it != null) {
onError?.invoke(it)
_onError.invoke(it)
} else {
onClose?.invoke()
_onClose.invoke()
}
}
}

17
src/jvmTest/kotlin/InMemoryTransport.kt
View file

@ -1,17 +1,14 @@
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
/**
* In-memory transport for creating clients and servers that talk to each other within the same process.
*/
class InMemoryTransport : Transport {
class InMemoryTransport : AbstractTransport() {
private var otherTransport: InMemoryTransport? = null
private val messageQueue: MutableList<JSONRPCMessage> = mutableListOf()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
/**
* Creates a pair of linked in-memory transports that can communicate with each other.
* One should be passed to a Client and one to a Server.
@ -30,7 +27,7 @@ class InMemoryTransport : Transport {
// Process any messages that were queued before start was called
while (messageQueue.isNotEmpty()) {
messageQueue.removeFirstOrNull()?.let { message ->
onMessage?.invoke(message) // todo?
_onMessage.invoke(message) // todo?
}
}
}
@ -39,16 +36,12 @@ class InMemoryTransport : Transport {
val other = otherTransport
otherTransport = null
other?.close()
onClose?.invoke()
_onClose.invoke()
}
override suspend fun send(message: JSONRPCMessage) {
val other = otherTransport ?: throw IllegalStateException("Not connected")
if (other.onMessage != null) {
other.onMessage?.invoke(message) // todo?
} else {
other.messageQueue.add(message)
}
other._onMessage.invoke(message)
}
}

8
src/jvmTest/kotlin/client/BaseTransportTest.kt
View file

@ -13,12 +13,12 @@ import kotlin.test.fail
abstract class BaseTransportTest {
protected suspend fun testClientOpenClose(client: Transport) {
client.onError = { error ->
client.onError { error ->
fail("Unexpected error: $error")
}
var didClose = false
client.onClose = { didClose = true }
client.onClose { didClose = true }
client.start()
assertFalse(didClose, "Transport should not be closed immediately after start")
@ -28,7 +28,7 @@ abstract class BaseTransportTest {
}
protected suspend fun testClientRead(client: Transport) {
client.onError = { error ->
client.onError { error ->
error.printStackTrace()
fail("Unexpected error: $error")
}
@ -40,7 +40,7 @@ abstract class BaseTransportTest {
val readMessages = mutableListOf<JSONRPCMessage>()
val finished = CompletableDeferred<Unit>()
client.onMessage = { message ->
client.onMessage { message ->
readMessages.add(message)
if (message == messages.last()) {
finished.complete(Unit)

26
src/jvmTest/kotlin/client/ClientTest.kt
View file

@ -37,7 +37,7 @@ import io.modelcontextprotocol.kotlin.sdk.client.ClientOptions
import org.junit.jupiter.api.Test
import io.modelcontextprotocol.kotlin.sdk.server.Server
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import kotlin.coroutines.cancellation.CancellationException
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
@ -48,7 +48,7 @@ class ClientTest {
@Test
fun `should initialize with matching protocol version`() = runTest {
var initialied = false
val clientTransport = object : Transport {
val clientTransport = object : AbstractTransport() {
override suspend fun start() {}
override suspend fun send(message: JSONRPCMessage) {
@ -68,15 +68,11 @@ class ClientTest {
result = result
)
onMessage?.invoke(response)
_onMessage.invoke(response)
}
override suspend fun close() {
}
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null
}
val client = Client(
@ -98,7 +94,7 @@ class ClientTest {
@Test
fun `should initialize with supported older protocol version`() = runTest {
val OLD_VERSION = SUPPORTED_PROTOCOL_VERSIONS[1]
val clientTransport = object : Transport {
val clientTransport = object : AbstractTransport() {
override suspend fun start() {}
override suspend fun send(message: JSONRPCMessage) {
@ -118,15 +114,11 @@ class ClientTest {
id = message.id,
result = result
)
onMessage?.invoke(response)
_onMessage.invoke(response)
}
override suspend fun close() {
}
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null
}
val client = Client(
@ -151,7 +143,7 @@ class ClientTest {
@Test
fun `should reject unsupported protocol version`() = runTest {
var closed = false
val clientTransport = object : Transport {
val clientTransport = object : AbstractTransport() {
override suspend fun start() {}
override suspend fun send(message: JSONRPCMessage) {
@ -172,16 +164,12 @@ class ClientTest {
result = result
)
onMessage?.invoke(response)
_onMessage.invoke(response)
}
override suspend fun close() {
closed = true
}
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null
}
val client = Client(

12
src/jvmTest/kotlin/client/InMemoryTransportTest.kt
View file

@ -44,7 +44,7 @@ class InMemoryTransportTest {
val message = InitializedNotification()
var receivedMessage: JSONRPCMessage? = null
serverTransport.onMessage = { msg ->
serverTransport.onMessage { msg ->
receivedMessage = msg
}
@ -61,7 +61,7 @@ class InMemoryTransportTest {
.toJSON()
var receivedMessage: JSONRPCMessage? = null
clientTransport.onMessage = { msg ->
clientTransport.onMessage { msg ->
receivedMessage = msg
}
@ -76,11 +76,11 @@ class InMemoryTransportTest {
var clientClosed = false
var serverClosed = false
clientTransport.onClose = {
clientTransport.onClose {
clientClosed = true
}
serverTransport.onClose = {
serverTransport.onClose {
serverClosed = true
}
@ -112,7 +112,7 @@ class InMemoryTransportTest {
.toJSON()
var receivedMessage: JSONRPCMessage? = null
serverTransport.onMessage = { msg ->
serverTransport.onMessage { msg ->
receivedMessage = msg
}
@ -121,4 +121,4 @@ class InMemoryTransportTest {
assertEquals(message, receivedMessage)
}
}
}
}

31
src/jvmTest/kotlin/client/SseTransportTest.kt
View file

@ -6,10 +6,13 @@ import io.ktor.server.application.*
import io.ktor.server.cio.*
import io.ktor.server.engine.*
import io.ktor.server.routing.*
import io.ktor.server.sse.sse
import io.ktor.util.collections.ConcurrentMap
import kotlinx.coroutines.test.runTest
import mcpSse
import mcpSseTransport
import io.modelcontextprotocol.kotlin.sdk.client.mcpSseTransport
import io.modelcontextprotocol.kotlin.sdk.server.SseServerTransport
import io.modelcontextprotocol.kotlin.sdk.server.mcpPostEndpoint
import io.modelcontextprotocol.kotlin.sdk.server.mcpSseTransport
import org.junit.jupiter.api.Test
private const val PORT = 8080
@ -19,8 +22,15 @@ class SseTransportTest : BaseTransportTest() {
fun `should start then close cleanly`() = runTest {
val server = embeddedServer(CIO, port = PORT) {
install(io.ktor.server.sse.SSE)
val transports = ConcurrentMap<String, SseServerTransport>()
routing {
mcpSse()
sse {
mcpSseTransport("", transports).start()
}
post {
mcpPostEndpoint(transports)
}
}
}.start(wait = false)
@ -42,12 +52,21 @@ class SseTransportTest : BaseTransportTest() {
fun `should read messages`() = runTest {
val server = embeddedServer(CIO, port = PORT) {
install(io.ktor.server.sse.SSE)
val transports = ConcurrentMap<String, SseServerTransport>()
routing {
mcpSseTransport {
onMessage = {
send(it)
sse {
mcpSseTransport("", transports).apply {
onMessage {
send(it)
}
start()
}
}
post {
mcpPostEndpoint(transports)
}
}
}.start(wait = false)

2
src/jvmTest/kotlin/client/WebSocketTransportTest.kt
View file

@ -33,7 +33,7 @@ class WebSocketTransportTest : BaseTransportTest() {
install(WebSockets)
routing {
mcpWebSocketTransport {
onMessage = {
onMessage {
send(it)
}

14
src/jvmTest/kotlin/server/StdioServerTransportTest.kt
View file

@ -54,12 +54,12 @@ class StdioServerTransportTest {
fun `should start then close cleanly`() {
runBlocking {
val server = StdioServerTransport(bufferedInput, printOutput)
server.onError = { error ->
server.onError { error ->
throw error
}
var didClose = false
server.onClose = {
server.onClose {
didClose = true
}
@ -75,14 +75,14 @@ class StdioServerTransportTest {
fun `should not read until started`() {
runBlocking {
val server = StdioServerTransport(bufferedInput, printOutput)
server.onError = { error ->
server.onError { error ->
throw error
}
var didRead = false
val readMessage = CompletableDeferred<JSONRPCMessage>()
server.onMessage = { message ->
server.onMessage { message ->
didRead = true
readMessage.complete(message)
}
@ -106,7 +106,7 @@ class StdioServerTransportTest {
fun `should read multiple messages`() {
runBlocking {
val server = StdioServerTransport(bufferedInput, printOutput)
server.onError = { error ->
server.onError { error ->
throw error
}
@ -118,7 +118,7 @@ class StdioServerTransportTest {
val readMessages = mutableListOf<JSONRPCMessage>()
val finished = CompletableDeferred<Unit>()
server.onMessage = { message ->
server.onMessage { message ->
readMessages.add(message)
if (message == messages[1]) {
finished.complete(Unit)
@ -141,4 +141,4 @@ class StdioServerTransportTest {
fun PipedOutputStream.write(s: String) {
write(s.toByteArray())
}
}

Some files were not shown because too many files have changed in this diff Show more