Migrate project to Multiplatform (KMP) (#22)
* Migrate to Kotlin multiplatform * Fix publication * Switch to using `alias` for the AtomicFu plugin --------- Co-authored-by: Alexander Sysoev <Alexander.Sysoev@jetbrains.com>
This commit is contained in:
Pavel Gorgulov
GitHub
parent
67d1f7775a
commit
e8c2f007a5
12 changed files with 160 additions and 125 deletions
54
src/jvmTest/kotlin/InMemoryTransport.kt
Normal file
54
src/jvmTest/kotlin/InMemoryTransport.kt
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
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 {
|
||||
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.
|
||||
*/
|
||||
companion object {
|
||||
fun createLinkedPair(): Pair<InMemoryTransport, InMemoryTransport> {
|
||||
val clientTransport = InMemoryTransport()
|
||||
val serverTransport = InMemoryTransport()
|
||||
clientTransport.otherTransport = serverTransport
|
||||
serverTransport.otherTransport = clientTransport
|
||||
return Pair(clientTransport, serverTransport)
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun start() {
|
||||
// Process any messages that were queued before start was called
|
||||
while (messageQueue.isNotEmpty()) {
|
||||
messageQueue.removeFirstOrNull()?.let { message ->
|
||||
onMessage?.invoke(message) // todo?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
val other = otherTransport
|
||||
otherTransport = null
|
||||
other?.close()
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
62
src/jvmTest/kotlin/client/BaseTransportTest.kt
Normal file
62
src/jvmTest/kotlin/client/BaseTransportTest.kt
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
package client
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.PingRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
|
||||
import io.modelcontextprotocol.kotlin.sdk.toJSON
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertFalse
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.test.fail
|
||||
|
||||
abstract class BaseTransportTest {
|
||||
protected suspend fun testClientOpenClose(client: Transport) {
|
||||
client.onError = { error ->
|
||||
fail("Unexpected error: $error")
|
||||
}
|
||||
|
||||
var didClose = false
|
||||
client.onClose = { didClose = true }
|
||||
|
||||
client.start()
|
||||
assertFalse(didClose, "Transport should not be closed immediately after start")
|
||||
|
||||
client.close()
|
||||
assertTrue(didClose, "Transport should be closed after close() call")
|
||||
}
|
||||
|
||||
protected suspend fun testClientRead(client: Transport) {
|
||||
client.onError = { error ->
|
||||
error.printStackTrace()
|
||||
fail("Unexpected error: $error")
|
||||
}
|
||||
|
||||
val messages = listOf<JSONRPCMessage>(
|
||||
PingRequest().toJSON(), InitializedNotification().toJSON()
|
||||
)
|
||||
|
||||
val readMessages = mutableListOf<JSONRPCMessage>()
|
||||
val finished = CompletableDeferred<Unit>()
|
||||
|
||||
client.onMessage = { message ->
|
||||
readMessages.add(message)
|
||||
if (message == messages.last()) {
|
||||
finished.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
client.start()
|
||||
|
||||
for (message in messages) {
|
||||
client.send(message)
|
||||
}
|
||||
|
||||
finished.await()
|
||||
|
||||
assertEquals(messages, readMessages, "Assert messages received")
|
||||
|
||||
client.close()
|
||||
}
|
||||
}
|
||||
45
src/jvmTest/kotlin/client/ClientIntegrationTest.kt
Normal file
45
src/jvmTest/kotlin/client/ClientIntegrationTest.kt
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
package client
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListToolsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.Client
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.io.asSink
|
||||
import kotlinx.io.asSource
|
||||
import kotlinx.io.buffered
|
||||
import org.junit.jupiter.api.Disabled
|
||||
import org.junit.jupiter.api.Test
|
||||
import java.net.Socket
|
||||
|
||||
class ClientIntegrationTest {
|
||||
|
||||
fun createTransport(): StdioClientTransport {
|
||||
val socket = Socket("localhost", 3000)
|
||||
|
||||
return StdioClientTransport(
|
||||
socket.inputStream.asSource().buffered(),
|
||||
socket.outputStream.asSink().buffered()
|
||||
)
|
||||
}
|
||||
|
||||
@Disabled("This test requires a running server")
|
||||
@Test
|
||||
fun testRequestTools() = runTest {
|
||||
val client = Client(
|
||||
Implementation("test", "1.0"),
|
||||
)
|
||||
|
||||
val transport = createTransport()
|
||||
try {
|
||||
client.connect(transport)
|
||||
|
||||
val response: ListToolsResult? = client.listTools()
|
||||
println(response?.tools)
|
||||
|
||||
} finally {
|
||||
transport.close()
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
510
src/jvmTest/kotlin/client/ClientTest.kt
Normal file
510
src/jvmTest/kotlin/client/ClientTest.kt
Normal file
|
|
@ -0,0 +1,510 @@
|
|||
package client
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.ClientCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateMessageRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateMessageResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.EmptyJsonObject
|
||||
import io.modelcontextprotocol.kotlin.sdk.Implementation
|
||||
import InMemoryTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializeRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializeResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCResponse
|
||||
import io.modelcontextprotocol.kotlin.sdk.LATEST_PROTOCOL_VERSION
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListResourcesRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListResourcesResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListRootsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListToolsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListToolsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.LoggingLevel
|
||||
import io.modelcontextprotocol.kotlin.sdk.LoggingMessageNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.Role
|
||||
import io.modelcontextprotocol.kotlin.sdk.SUPPORTED_PROTOCOL_VERSIONS
|
||||
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.TextContent
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.TimeoutCancellationException
|
||||
import kotlinx.coroutines.cancel
|
||||
import kotlinx.coroutines.joinAll
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.serialization.json.buildJsonObject
|
||||
import kotlinx.serialization.json.put
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.Client
|
||||
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 kotlin.coroutines.cancellation.CancellationException
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertFailsWith
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.test.fail
|
||||
|
||||
class ClientTest {
|
||||
@Test
|
||||
fun `should initialize with matching protocol version`() = runTest {
|
||||
var initialied = false
|
||||
val clientTransport = object : Transport {
|
||||
override suspend fun start() {}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage) {
|
||||
if (message !is JSONRPCRequest) return
|
||||
initialied = true
|
||||
val result = InitializeResult(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ServerCapabilities(),
|
||||
serverInfo = Implementation(
|
||||
name = "test",
|
||||
version = "1.0"
|
||||
)
|
||||
)
|
||||
|
||||
val response = JSONRPCResponse(
|
||||
id = message.id,
|
||||
result = result
|
||||
)
|
||||
|
||||
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(
|
||||
clientInfo = Implementation(
|
||||
name = "test client",
|
||||
version = "1.0"
|
||||
),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(
|
||||
sampling = EmptyJsonObject
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
client.connect(clientTransport)
|
||||
assertTrue(initialied)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should initialize with supported older protocol version`() = runTest {
|
||||
val OLD_VERSION = SUPPORTED_PROTOCOL_VERSIONS[1]
|
||||
val clientTransport = object : Transport {
|
||||
override suspend fun start() {}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage) {
|
||||
if (message !is JSONRPCRequest) return
|
||||
check(message.method == Method.Defined.Initialize.value)
|
||||
|
||||
val result = InitializeResult(
|
||||
protocolVersion = OLD_VERSION,
|
||||
capabilities = ServerCapabilities(),
|
||||
serverInfo = Implementation(
|
||||
name = "test",
|
||||
version = "1.0"
|
||||
)
|
||||
)
|
||||
|
||||
val response = JSONRPCResponse(
|
||||
id = message.id,
|
||||
result = result
|
||||
)
|
||||
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(
|
||||
clientInfo = Implementation(
|
||||
name = "test client",
|
||||
version = "1.0"
|
||||
),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(
|
||||
sampling = EmptyJsonObject
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
client.connect(clientTransport)
|
||||
assertEquals(
|
||||
Implementation("test", "1.0"),
|
||||
client.serverVersion
|
||||
)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should reject unsupported protocol version`() = runTest {
|
||||
var closed = false
|
||||
val clientTransport = object : Transport {
|
||||
override suspend fun start() {}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage) {
|
||||
if (message !is JSONRPCRequest) return
|
||||
check(message.method == Method.Defined.Initialize.value)
|
||||
|
||||
val result = InitializeResult(
|
||||
protocolVersion = "invalid-version",
|
||||
capabilities = ServerCapabilities(),
|
||||
serverInfo = Implementation(
|
||||
name = "test",
|
||||
version = "1.0"
|
||||
)
|
||||
)
|
||||
|
||||
val response = JSONRPCResponse(
|
||||
id = message.id,
|
||||
result = result
|
||||
)
|
||||
|
||||
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(
|
||||
clientInfo = Implementation(
|
||||
name = "test client",
|
||||
version = "1.0"
|
||||
),
|
||||
options = ClientOptions()
|
||||
)
|
||||
|
||||
assertFailsWith<IllegalStateException>("Server's protocol version is not supported: invalid-version") {
|
||||
client.connect(clientTransport)
|
||||
}
|
||||
|
||||
assertTrue(closed)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should respect server capabilities`() = runTest {
|
||||
val serverOptions = ServerOptions(
|
||||
capabilities = ServerCapabilities(
|
||||
resources = ServerCapabilities.Resources(null, null),
|
||||
tools = ServerCapabilities.Tools(null)
|
||||
)
|
||||
)
|
||||
val server = Server(
|
||||
Implementation(name = "test server", version = "1.0"),
|
||||
serverOptions
|
||||
)
|
||||
|
||||
server.setRequestHandler<InitializeRequest>(Method.Defined.Initialize) { request, _ ->
|
||||
InitializeResult(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ServerCapabilities(
|
||||
resources = ServerCapabilities.Resources(null, null),
|
||||
tools = ServerCapabilities.Tools(null)
|
||||
),
|
||||
serverInfo = Implementation(name = "test", version = "1.0")
|
||||
)
|
||||
}
|
||||
|
||||
server.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { request, _ ->
|
||||
ListResourcesResult(resources = emptyList(), nextCursor = null)
|
||||
}
|
||||
|
||||
server.setRequestHandler<ListToolsRequest>(Method.Defined.ToolsList) { request, _ ->
|
||||
ListToolsResult(tools = emptyList(), nextCursor = null)
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(sampling = EmptyJsonObject),
|
||||
)
|
||||
)
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
}
|
||||
).joinAll()
|
||||
|
||||
// Server supports resources and tools, but not prompts
|
||||
val caps = client.serverCapabilities
|
||||
assertEquals(ServerCapabilities.Resources(null, null), caps?.resources)
|
||||
assertEquals(ServerCapabilities.Tools(null), caps?.tools)
|
||||
assertTrue(caps?.prompts == null) // or check that prompts are absent
|
||||
|
||||
// These should not throw
|
||||
client.listResources()
|
||||
client.listTools()
|
||||
|
||||
// This should fail because prompts are not supported
|
||||
val ex = assertFailsWith<IllegalStateException> {
|
||||
client.listPrompts()
|
||||
}
|
||||
assertTrue(ex.message?.contains("Server does not support prompts") == true)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should respect client notification capabilities`() = runTest {
|
||||
val server = Server(
|
||||
Implementation(name = "test server", version = "1.0"),
|
||||
ServerOptions(capabilities = ServerCapabilities())
|
||||
)
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(
|
||||
roots = ClientCapabilities.Roots(listChanged = true)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
println("Server connected")
|
||||
}
|
||||
).joinAll()
|
||||
|
||||
// This should not throw because the client supports roots.listChanged
|
||||
client.sendRootsListChanged()
|
||||
|
||||
// Create a new client without the roots.listChanged capability
|
||||
val clientWithoutCapability = Client(
|
||||
clientInfo = Implementation(name = "test client without capability", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(),
|
||||
// enforceStrictCapabilities = true // TODO()
|
||||
)
|
||||
)
|
||||
|
||||
clientWithoutCapability.connect(clientTransport)
|
||||
// Using the same transport pair might not be realistic - in a real scenario you'd create another pair.
|
||||
// Adjust if necessary.
|
||||
|
||||
// This should fail
|
||||
val ex = assertFailsWith<IllegalStateException> {
|
||||
clientWithoutCapability.sendRootsListChanged()
|
||||
}
|
||||
assertTrue(ex.message?.startsWith("Client does not support") == true)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should respect server notification capabilities`() = runTest {
|
||||
val server = Server(
|
||||
Implementation(name = "test server", version = "1.0"),
|
||||
ServerOptions(
|
||||
capabilities = ServerCapabilities(
|
||||
logging = EmptyJsonObject,
|
||||
resources = ServerCapabilities.Resources(listChanged = true, subscribe = null)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities()
|
||||
)
|
||||
)
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
println("Server connected")
|
||||
}
|
||||
).joinAll()
|
||||
|
||||
// These should not throw
|
||||
val jsonObject = buildJsonObject {
|
||||
put("name", "John")
|
||||
put("age", 30)
|
||||
put("isStudent", false)
|
||||
}
|
||||
server.sendLoggingMessage(
|
||||
LoggingMessageNotification(
|
||||
level = LoggingLevel.info,
|
||||
data = jsonObject
|
||||
)
|
||||
)
|
||||
server.sendResourceListChanged()
|
||||
|
||||
// This should fail because the server doesn't have the tools capability
|
||||
val ex = assertFailsWith<IllegalStateException> {
|
||||
server.sendToolListChanged()
|
||||
}
|
||||
assertTrue(ex.message?.contains("Server does not support notifying of tool list changes") == true)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should handle client cancelling a request`() = runTest {
|
||||
val server = Server(
|
||||
Implementation(name = "test server", version = "1.0"),
|
||||
ServerOptions(
|
||||
capabilities = ServerCapabilities(resources = ServerCapabilities.Resources(listChanged = null, subscribe = null))
|
||||
)
|
||||
)
|
||||
|
||||
val def = CompletableDeferred<Unit>()
|
||||
val defTimeOut = CompletableDeferred<Unit>()
|
||||
server.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, extra ->
|
||||
// Simulate delay
|
||||
def.complete(Unit)
|
||||
try {
|
||||
kotlinx.coroutines.delay(1000)
|
||||
} catch (e: CancellationException) {
|
||||
defTimeOut.complete(Unit)
|
||||
throw e
|
||||
}
|
||||
fail("Shouldn't have been called")
|
||||
ListResourcesResult(resources = emptyList())
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(capabilities = ClientCapabilities())
|
||||
)
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
println("Server connected")
|
||||
}
|
||||
).joinAll()
|
||||
|
||||
val defCancel = CompletableDeferred<Unit>()
|
||||
val job = launch {
|
||||
try {
|
||||
client.listResources()
|
||||
} catch (e: CancellationException) {
|
||||
defCancel.complete(Unit)
|
||||
assertEquals("Cancelled by test", e.message)
|
||||
}
|
||||
}
|
||||
def.await()
|
||||
runCatching { job.cancel("Cancelled by test") }
|
||||
defCancel.await()
|
||||
defTimeOut.await()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should handle request timeout`() = runTest {
|
||||
val server = Server(
|
||||
Implementation(name = "test server", version = "1.0"),
|
||||
ServerOptions(
|
||||
capabilities = ServerCapabilities(resources = ServerCapabilities.Resources(listChanged = null, subscribe = null))
|
||||
)
|
||||
)
|
||||
|
||||
server.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, extra ->
|
||||
// Simulate a delayed response
|
||||
// Wait ~100ms unless cancelled
|
||||
try {
|
||||
kotlinx.coroutines.withTimeout(100L) {
|
||||
// Just delay here, if timeout is 0 on client side this won't return in time
|
||||
kotlinx.coroutines.delay(100)
|
||||
}
|
||||
} catch (_: Exception) {
|
||||
// If aborted, just rethrow or return early
|
||||
}
|
||||
ListResourcesResult(resources = emptyList())
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(capabilities = ClientCapabilities())
|
||||
)
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
println("Server connected")
|
||||
}
|
||||
).joinAll()
|
||||
|
||||
// Request with 1 msec timeout should fail immediately
|
||||
val ex = assertFailsWith<Exception> {
|
||||
kotlinx.coroutines.withTimeout(1) {
|
||||
client.listResources()
|
||||
}
|
||||
}
|
||||
assertTrue(ex is TimeoutCancellationException)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should only allow setRequestHandler for declared capabilities`() = runTest {
|
||||
val client = Client(
|
||||
clientInfo = Implementation(
|
||||
name = "test client",
|
||||
version = "1.0"
|
||||
),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(
|
||||
sampling = EmptyJsonObject
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
client.setRequestHandler<CreateMessageRequest>(Method.Defined.SamplingCreateMessage) { request, _ ->
|
||||
CreateMessageResult(
|
||||
model = "test-model",
|
||||
role = Role.assistant,
|
||||
content = TextContent(
|
||||
text = "Test response"
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
assertFailsWith<IllegalStateException>("Client does not support roots capability (required for RootsList)") {
|
||||
client.setRequestHandler<ListRootsRequest>(Method.Defined.RootsList) { _, _ -> null }
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
124
src/jvmTest/kotlin/client/InMemoryTransportTest.kt
Normal file
124
src/jvmTest/kotlin/client/InMemoryTransportTest.kt
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
package client
|
||||
|
||||
import InMemoryTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import io.modelcontextprotocol.kotlin.sdk.toJSON
|
||||
import org.junit.jupiter.api.BeforeEach
|
||||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.api.assertThrows
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertTrue
|
||||
|
||||
class InMemoryTransportTest {
|
||||
private lateinit var clientTransport: InMemoryTransport
|
||||
private lateinit var serverTransport: InMemoryTransport
|
||||
|
||||
@BeforeEach
|
||||
fun setUp() {
|
||||
val (client, server) = InMemoryTransport.createLinkedPair()
|
||||
clientTransport = client
|
||||
serverTransport = server
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should create linked pair`() {
|
||||
assertNotNull(clientTransport)
|
||||
assertNotNull(serverTransport)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should start without error`() {
|
||||
runBlocking {
|
||||
clientTransport.start()
|
||||
serverTransport.start()
|
||||
// If no exception is thrown, the test passes
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should send message from client to server`() {
|
||||
runBlocking {
|
||||
val message = InitializedNotification()
|
||||
|
||||
var receivedMessage: JSONRPCMessage? = null
|
||||
serverTransport.onMessage = { msg ->
|
||||
receivedMessage = msg
|
||||
}
|
||||
|
||||
val rpcNotification = message.toJSON()
|
||||
clientTransport.send(rpcNotification)
|
||||
assertEquals(rpcNotification, receivedMessage)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should send message from server to client`() {
|
||||
runBlocking {
|
||||
val message = InitializedNotification()
|
||||
.toJSON()
|
||||
|
||||
var receivedMessage: JSONRPCMessage? = null
|
||||
clientTransport.onMessage = { msg ->
|
||||
receivedMessage = msg
|
||||
}
|
||||
|
||||
serverTransport.send(message)
|
||||
assertEquals(message, receivedMessage)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should handle close`() {
|
||||
runBlocking {
|
||||
var clientClosed = false
|
||||
var serverClosed = false
|
||||
|
||||
clientTransport.onClose = {
|
||||
clientClosed = true
|
||||
}
|
||||
|
||||
serverTransport.onClose = {
|
||||
serverClosed = true
|
||||
}
|
||||
|
||||
clientTransport.close()
|
||||
assertTrue(clientClosed)
|
||||
assertTrue(serverClosed)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should throw error when sending after close`() {
|
||||
runBlocking {
|
||||
clientTransport.close()
|
||||
|
||||
assertThrows<IllegalStateException> {
|
||||
runBlocking {
|
||||
clientTransport.send(
|
||||
InitializedNotification().toJSON()
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should queue messages sent before start`() {
|
||||
runBlocking {
|
||||
val message = InitializedNotification()
|
||||
.toJSON()
|
||||
|
||||
var receivedMessage: JSONRPCMessage? = null
|
||||
serverTransport.onMessage = { msg ->
|
||||
receivedMessage = msg
|
||||
}
|
||||
|
||||
clientTransport.send(message)
|
||||
serverTransport.start()
|
||||
assertEquals(message, receivedMessage)
|
||||
}
|
||||
}
|
||||
}
|
||||
66
src/jvmTest/kotlin/client/SseTransportTest.kt
Normal file
66
src/jvmTest/kotlin/client/SseTransportTest.kt
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
package client
|
||||
|
||||
import io.ktor.client.*
|
||||
import io.ktor.client.plugins.sse.*
|
||||
import io.ktor.server.application.*
|
||||
import io.ktor.server.cio.*
|
||||
import io.ktor.server.engine.*
|
||||
import io.ktor.server.routing.*
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import mcpSse
|
||||
import mcpSseTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.mcpSseTransport
|
||||
import org.junit.jupiter.api.Test
|
||||
|
||||
private const val PORT = 8080
|
||||
|
||||
class SseTransportTest : BaseTransportTest() {
|
||||
@Test
|
||||
fun `should start then close cleanly`() = runTest {
|
||||
val server = embeddedServer(CIO, port = PORT) {
|
||||
install(io.ktor.server.sse.SSE)
|
||||
routing {
|
||||
mcpSse()
|
||||
}
|
||||
}.start(wait = false)
|
||||
|
||||
val client = HttpClient {
|
||||
install(SSE)
|
||||
}.mcpSseTransport {
|
||||
url {
|
||||
host = "localhost"
|
||||
port = PORT
|
||||
}
|
||||
}
|
||||
|
||||
testClientOpenClose(client)
|
||||
|
||||
server.stop()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should read messages`() = runTest {
|
||||
val server = embeddedServer(CIO, port = PORT) {
|
||||
install(io.ktor.server.sse.SSE)
|
||||
routing {
|
||||
mcpSseTransport {
|
||||
onMessage = {
|
||||
send(it)
|
||||
}
|
||||
}
|
||||
}
|
||||
}.start(wait = false)
|
||||
|
||||
val client = HttpClient {
|
||||
install(SSE)
|
||||
}.mcpSseTransport {
|
||||
url {
|
||||
host = "localhost"
|
||||
port = PORT
|
||||
}
|
||||
}
|
||||
|
||||
testClientRead(client)
|
||||
server.stop()
|
||||
}
|
||||
}
|
||||
48
src/jvmTest/kotlin/client/StdioClientTransportTest.kt
Normal file
48
src/jvmTest/kotlin/client/StdioClientTransportTest.kt
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
package client
|
||||
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport
|
||||
import kotlinx.io.asSink
|
||||
import kotlinx.io.asSource
|
||||
import kotlinx.io.buffered
|
||||
import org.junit.jupiter.api.Test
|
||||
|
||||
class StdioClientTransportTest : BaseTransportTest() {
|
||||
@Test
|
||||
fun `should start then close cleanly`() = runTest {
|
||||
// Run process "/usr/bin/tee"
|
||||
val processBuilder = ProcessBuilder("/usr/bin/tee")
|
||||
val process = processBuilder.start()
|
||||
|
||||
val input = process.inputStream.asSource().buffered()
|
||||
val output = process.outputStream.asSink().buffered()
|
||||
|
||||
val client = StdioClientTransport(
|
||||
input = input,
|
||||
output = output
|
||||
)
|
||||
|
||||
testClientOpenClose(client)
|
||||
|
||||
process.destroy()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should read messages`() = runTest {
|
||||
val processBuilder = ProcessBuilder("/usr/bin/tee")
|
||||
val process = processBuilder.start()
|
||||
|
||||
val input = process.inputStream.asSource().buffered()
|
||||
val output = process.outputStream.asSink().buffered()
|
||||
|
||||
val client = StdioClientTransport(
|
||||
input = input,
|
||||
output = output
|
||||
)
|
||||
|
||||
testClientRead(client)
|
||||
|
||||
process.waitFor()
|
||||
process.destroy()
|
||||
}
|
||||
}
|
||||
47
src/jvmTest/kotlin/client/TypesTest.kt
Normal file
47
src/jvmTest/kotlin/client/TypesTest.kt
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
package client
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.Request
|
||||
import org.junit.jupiter.api.Test
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.McpJson
|
||||
|
||||
class TypesTest {
|
||||
|
||||
@Test
|
||||
fun testRequestResult() {
|
||||
val message = "{\"result\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{\"tools\":{\"listChanged\":true},\"resources\":{}},\"serverInfo\":{\"name\":\"jetbrains/proxy\",\"version\":\"0.1.0\"}},\"jsonrpc\":\"2.0\",\"id\":1}"
|
||||
McpJson.decodeFromString<JSONRPCMessage>(message)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testRequestError() {
|
||||
val message = "{\"method\":\"initialize\", \"protocolVersion\":\"2024-11-05\",\"capabilities\":{\"experimental\":{},\"sampling\":{}},\"clientInfo\":{\"name\":\"test client\",\"version\":\"1.0\"},\"_meta\":{}}"
|
||||
McpJson.decodeFromString<Request>(message)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testJSONRPCMessage() {
|
||||
val line = "{\"result\":{\"content\":[{\"type\":\"text\"}],\"isError\":false},\"jsonrpc\":\"2.0\",\"id\":4}"
|
||||
McpJson.decodeFromString<JSONRPCMessage>(line)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testJSONRPCMessageWithStringId() {
|
||||
val line = """
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "initialize",
|
||||
"id": "ebf9f64a-0",
|
||||
"params": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {},
|
||||
"clientInfo": {
|
||||
"name": "mcp-java-client",
|
||||
"version": "0.2.0"
|
||||
}
|
||||
}
|
||||
}
|
||||
""".trimIndent()
|
||||
McpJson.decodeFromString<JSONRPCMessage>(line)
|
||||
}
|
||||
}
|
||||
52
src/jvmTest/kotlin/client/WebSocketTransportTest.kt
Normal file
52
src/jvmTest/kotlin/client/WebSocketTransportTest.kt
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
package client
|
||||
|
||||
import io.ktor.server.testing.*
|
||||
import io.ktor.server.websocket.*
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.mcpWebSocketTransport
|
||||
import org.junit.jupiter.api.Test
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.mcpWebSocket
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.mcpWebSocketTransport
|
||||
import org.junit.jupiter.api.Disabled
|
||||
|
||||
class WebSocketTransportTest : BaseTransportTest() {
|
||||
@Test
|
||||
@Disabled("Test disabled for investigation #17")
|
||||
fun `should start then close cleanly`() = testApplication {
|
||||
install(WebSockets)
|
||||
routing {
|
||||
mcpWebSocket()
|
||||
}
|
||||
|
||||
val client = createClient {
|
||||
install(io.ktor.client.plugins.websocket.WebSockets)
|
||||
}.mcpWebSocketTransport()
|
||||
|
||||
testClientOpenClose(client)
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled("Test disabled for investigation #17")
|
||||
fun `should read messages`() = testApplication {
|
||||
val clientFinished = CompletableDeferred<Unit>()
|
||||
|
||||
install(WebSockets)
|
||||
routing {
|
||||
mcpWebSocketTransport {
|
||||
onMessage = {
|
||||
send(it)
|
||||
}
|
||||
|
||||
clientFinished.await()
|
||||
}
|
||||
}
|
||||
|
||||
val client = createClient {
|
||||
install(io.ktor.client.plugins.websocket.WebSockets)
|
||||
}.mcpWebSocketTransport()
|
||||
|
||||
testClientRead(client)
|
||||
|
||||
clientFinished.complete(Unit)
|
||||
}
|
||||
}
|
||||
144
src/jvmTest/kotlin/server/StdioServerTransportTest.kt
Normal file
144
src/jvmTest/kotlin/server/StdioServerTransportTest.kt
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
package server
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.PingRequest
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.StdioServerTransport
|
||||
import org.junit.jupiter.api.Assertions.*
|
||||
import org.junit.jupiter.api.BeforeEach
|
||||
import org.junit.jupiter.api.Test
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.toJSON
|
||||
import kotlinx.io.Sink
|
||||
import kotlinx.io.Source
|
||||
import kotlinx.io.asSink
|
||||
import kotlinx.io.asSource
|
||||
import kotlinx.io.buffered
|
||||
import java.io.*
|
||||
|
||||
class StdioServerTransportTest {
|
||||
private lateinit var input: PipedInputStream
|
||||
private lateinit var inputWriter: PipedOutputStream
|
||||
private lateinit var outputBuffer: ReadBuffer
|
||||
private lateinit var output: ByteArrayOutputStream
|
||||
|
||||
// We'll store the wrapped streams that meet the constructor requirements
|
||||
private lateinit var bufferedInput: Source
|
||||
private lateinit var printOutput: Sink
|
||||
|
||||
@BeforeEach
|
||||
fun setUp() {
|
||||
// Simulate an input stream that we can push data into using inputWriter.
|
||||
input = PipedInputStream()
|
||||
inputWriter = PipedOutputStream(input)
|
||||
|
||||
outputBuffer = ReadBuffer()
|
||||
|
||||
// A custom ByteArrayOutputStream that appends all written data into outputBuffer
|
||||
output = object : ByteArrayOutputStream() {
|
||||
override fun write(b: ByteArray, off: Int, len: Int) {
|
||||
super.write(b, off, len)
|
||||
outputBuffer.append(b.copyOfRange(off, off + len))
|
||||
}
|
||||
}
|
||||
|
||||
bufferedInput = input.asSource().buffered()
|
||||
|
||||
printOutput = output.asSink().buffered()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should start then close cleanly`() {
|
||||
runBlocking {
|
||||
val server = StdioServerTransport(bufferedInput, printOutput)
|
||||
server.onError = { error ->
|
||||
throw error
|
||||
}
|
||||
|
||||
var didClose = false
|
||||
server.onClose = {
|
||||
didClose = true
|
||||
}
|
||||
|
||||
server.start()
|
||||
assertFalse(didClose, "Should not have closed yet")
|
||||
|
||||
server.close()
|
||||
assertTrue(didClose, "Should have closed after calling close()")
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should not read until started`() {
|
||||
runBlocking {
|
||||
val server = StdioServerTransport(bufferedInput, printOutput)
|
||||
server.onError = { error ->
|
||||
throw error
|
||||
}
|
||||
|
||||
var didRead = false
|
||||
val readMessage = CompletableDeferred<JSONRPCMessage>()
|
||||
|
||||
server.onMessage = { message ->
|
||||
didRead = true
|
||||
readMessage.complete(message)
|
||||
}
|
||||
|
||||
val message = PingRequest().toJSON()
|
||||
|
||||
// Push message before the server started
|
||||
val serialized = serializeMessage(message)
|
||||
inputWriter.write(serialized)
|
||||
inputWriter.flush()
|
||||
|
||||
assertFalse(didRead, "Should not have read message before start")
|
||||
|
||||
server.start()
|
||||
val received = readMessage.await()
|
||||
assertEquals(message, received)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should read multiple messages`() {
|
||||
runBlocking {
|
||||
val server = StdioServerTransport(bufferedInput, printOutput)
|
||||
server.onError = { error ->
|
||||
throw error
|
||||
}
|
||||
|
||||
val messages = listOf(
|
||||
PingRequest().toJSON(),
|
||||
InitializedNotification().toJSON(),
|
||||
)
|
||||
|
||||
val readMessages = mutableListOf<JSONRPCMessage>()
|
||||
val finished = CompletableDeferred<Unit>()
|
||||
|
||||
server.onMessage = { message ->
|
||||
readMessages.add(message)
|
||||
if (message == messages[1]) {
|
||||
finished.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
// Push both messages before starting the server
|
||||
for (m in messages) {
|
||||
inputWriter.write(serializeMessage(m))
|
||||
}
|
||||
inputWriter.flush()
|
||||
|
||||
server.start()
|
||||
finished.await()
|
||||
|
||||
assertEquals(messages, readMessages)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fun PipedOutputStream.write(s: String) {
|
||||
write(s.toByteArray())
|
||||
}
|
||||
58
src/jvmTest/kotlin/shared/ReadBufferTest.kt
Normal file
58
src/jvmTest/kotlin/shared/ReadBufferTest.kt
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
package shared
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCNotification
|
||||
import kotlinx.serialization.encodeToString
|
||||
import kotlinx.serialization.json.Json
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
|
||||
import org.junit.jupiter.api.Assertions.assertEquals
|
||||
import org.junit.jupiter.api.Assertions.assertNull
|
||||
import org.junit.jupiter.api.Test
|
||||
import java.nio.charset.StandardCharsets
|
||||
|
||||
class ReadBufferTest {
|
||||
private val testMessage: JSONRPCMessage = JSONRPCNotification(method = "foobar")
|
||||
|
||||
private val json = Json {
|
||||
ignoreUnknownKeys = true
|
||||
encodeDefaults = true
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should have no messages after initialization`() {
|
||||
val readBuffer = ReadBuffer()
|
||||
assertNull(readBuffer.readMessage())
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should only yield a message after a newline`() {
|
||||
val readBuffer = ReadBuffer()
|
||||
|
||||
// Append message without newline
|
||||
val messageBytes = json.encodeToString(testMessage)
|
||||
.toByteArray(StandardCharsets.UTF_8)
|
||||
readBuffer.append(messageBytes)
|
||||
assertNull(readBuffer.readMessage())
|
||||
|
||||
// Append newline and verify message is now available
|
||||
readBuffer.append("\n".toByteArray(StandardCharsets.UTF_8))
|
||||
assertEquals(testMessage, readBuffer.readMessage())
|
||||
assertNull(readBuffer.readMessage())
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should be reusable after clearing`() {
|
||||
val readBuffer = ReadBuffer()
|
||||
|
||||
readBuffer.append("foobar".toByteArray(Charsets.UTF_8))
|
||||
readBuffer.clear()
|
||||
assertNull(readBuffer.readMessage())
|
||||
|
||||
val messageJson = serializeMessage(testMessage)
|
||||
readBuffer.append(messageJson.toByteArray(Charsets.UTF_8))
|
||||
readBuffer.append("\n".toByteArray(Charsets.UTF_8))
|
||||
val message = readBuffer.readMessage()
|
||||
assertEquals(testMessage, message)
|
||||
}
|
||||
}
|
||||
162
src/jvmTest/kotlin/sse.ktor.kt
Normal file
162
src/jvmTest/kotlin/sse.ktor.kt
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
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
|
||||
}
|
||||
Loading…
Reference in a new issue