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

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
8 changed files with 296 additions and 355 deletions
Download patch file Download diff file Expand all files Collapse all files

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())
}
}

162
src/jvmTest/kotlin/sse.ktor.kt
View file

@ -1,162 +0,0 @@
import io.ktor.http.HttpStatusCode
import io.ktor.server.application.ApplicationCall
import io.ktor.server.response.respondText
import io.ktor.server.routing.Route
import io.ktor.server.routing.RoutingContext
import io.ktor.server.routing.application
import io.ktor.server.routing.post
import io.ktor.server.sse.ServerSSESession
import io.ktor.server.sse.sse
import io.ktor.util.AttributeKey
import io.ktor.util.Attributes
import io.modelcontextprotocol.kotlin.sdk.server.SSEServerTransport
import kotlinx.coroutines.CompletableDeferred
import io.modelcontextprotocol.kotlin.sdk.Implementation
import io.modelcontextprotocol.kotlin.sdk.LIB_VERSION
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
import io.modelcontextprotocol.kotlin.sdk.server.SESSION_ID_PARAM
import io.modelcontextprotocol.kotlin.sdk.server.Server
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
import io.modelcontextprotocol.kotlin.sdk.shared.IMPLEMENTATION_NAME
typealias IncomingHandler = suspend RoutingContext.(forward: suspend () -> Unit) -> Unit
fun Route.mcpSse(
options: ServerOptions? = null,
incomingPath: String = "",
incomingHandler: (IncomingHandler)? = null,
handler: suspend Server.() -> Unit = {},
) {
sse {
createMcpServer(this, incomingPath, options, handler)
}
setupPostRoute(incomingPath, incomingHandler)
}
fun Route.mcpSse(
path: String,
incomingPath: String = path,
options: ServerOptions? = null,
incomingHandler: (IncomingHandler)? = null,
handler: suspend Server.() -> Unit = {},
) {
sse(path) {
createMcpServer(this, incomingPath, options, handler)
}
setupPostRoute(incomingPath, incomingHandler)
}
fun Route.mcpSseTransport(
incomingPath: String = "",
incomingHandler: (IncomingHandler)? = null,
handler: suspend SSEServerTransport.() -> Unit = {},
) {
sse {
val transport = createMcpTransport(this, incomingPath)
handler(transport)
transport.start()
transport.close()
}
setupPostRoute(incomingPath, incomingHandler)
}
fun Route.mcpSseTransport(
path: String,
incomingPath: String = path,
incomingHandler: (IncomingHandler)? = null,
handler: suspend SSEServerTransport.() -> Unit = {},
) {
sse(path) {
val transport = createMcpTransport(this, incomingPath)
transport.start()
handler(transport)
transport.close()
}
setupPostRoute(incomingPath, incomingHandler)
}
internal val McpServersKey = AttributeKey<Attributes>("mcp-servers")
private fun String.asAttributeKey() = AttributeKey<SSEServerTransport>(this)
private suspend fun Route.forwardMcpMessage(call: ApplicationCall) {
val sessionId = call.request.queryParameters[SESSION_ID_PARAM]
?.asAttributeKey()
?: run {
call.sessionNotFound()
return
}
application.attributes.getOrNull(McpServersKey)
?.get(sessionId)
?.handlePostMessage(call)
?: call.sessionNotFound()
}
private suspend fun ApplicationCall.sessionNotFound() {
respondText("Session not found", status = HttpStatusCode.NotFound)
}
private fun Route.setupPostRoute(incomingPath: String, incomingHandler: IncomingHandler?) {
post(incomingPath) {
if (incomingHandler != null) {
incomingHandler {
forwardMcpMessage(call)
}
} else {
forwardMcpMessage(call)
}
}
}
private suspend fun Route.createMcpServer(
session: ServerSSESession,
incomingPath: String,
options: ServerOptions?,
handler: suspend Server.() -> Unit = {},
) {
val transport = createMcpTransport(session, incomingPath)
val closed = CompletableDeferred<Unit>()
val server = Server(
serverInfo = Implementation(
name = IMPLEMENTATION_NAME,
version = LIB_VERSION,
),
options = options ?: ServerOptions(
capabilities = ServerCapabilities(
prompts = ServerCapabilities.Prompts(listChanged = null),
resources = ServerCapabilities.Resources(subscribe = null, listChanged = null),
tools = ServerCapabilities.Tools(listChanged = null),
)
),
onCloseCallback = {
closed.complete(Unit)
},
)
server.connect(transport)
handler(server)
server.close()
}
private fun Route.createMcpTransport(
session: ServerSSESession,
incomingPath: String,
): SSEServerTransport {
val transport = SSEServerTransport(
endpoint = incomingPath,
session = session,
)
application.attributes
.computeIfAbsent(McpServersKey) { Attributes(concurrent = true) }
.put(transport.sessionId.asAttributeKey(), transport)
return transport
}