From c8a4236d5dead7204f10934bf3cbeb405a839c2d Mon Sep 17 00:00:00 2001 From: devcrocod Date: Wed, 26 Mar 2025 16:32:42 +0300 Subject: [PATCH] Refactor `onCloseCallback` to `onClose` and update related usages across the codebase --- .../kotlin-mcp-server/src/main/kotlin/Main.kt | 24 +++++++------------ .../sample/server/McpWeatherServer.kt | 2 +- .../kotlin/sdk/server/Server.kt | 2 -- src/jvmTest/kotlin/InMemoryTransport.kt | 1 - .../kotlin/client/InMemoryTransportTest.kt | 8 +++---- 5 files changed, 13 insertions(+), 24 deletions(-) diff --git a/samples/kotlin-mcp-server/src/main/kotlin/Main.kt b/samples/kotlin-mcp-server/src/main/kotlin/Main.kt index 98083574..a73b5b74 100644 --- a/samples/kotlin-mcp-server/src/main/kotlin/Main.kt +++ b/samples/kotlin-mcp-server/src/main/kotlin/Main.kt @@ -14,12 +14,11 @@ import io.modelcontextprotocol.kotlin.sdk.PromptMessage import io.modelcontextprotocol.kotlin.sdk.Role import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities import io.modelcontextprotocol.kotlin.sdk.Tool -import io.modelcontextprotocol.kotlin.sdk.server.MCP -import io.modelcontextprotocol.kotlin.sdk.server.SSEServerTransport import io.modelcontextprotocol.kotlin.sdk.server.Server import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions +import io.modelcontextprotocol.kotlin.sdk.server.SseServerTransport import io.modelcontextprotocol.kotlin.sdk.server.StdioServerTransport -import kotlinx.coroutines.CompletableDeferred +import io.modelcontextprotocol.kotlin.sdk.server.mcp import kotlinx.coroutines.Job import kotlinx.coroutines.runBlocking import kotlinx.io.asSink @@ -48,8 +47,6 @@ fun main(args: Array) { } fun configureServer(): Server { - val def = CompletableDeferred() - val server = Server( Implementation( name = "mcp-kotlin test server", @@ -61,10 +58,7 @@ fun configureServer(): Server { resources = ServerCapabilities.Resources(subscribe = true, listChanged = true), tools = ServerCapabilities.Tools(listChanged = true), ) - ), - onCloseCallback = { - def.complete(Unit) - } + ) ) server.addPrompt( @@ -129,7 +123,7 @@ fun runMcpServerUsingStdio() { runBlocking { server.connect(transport) val done = Job() - server.onCloseCallback = { + server.onClose { done.complete() } done.join() @@ -146,7 +140,7 @@ fun runSseMcpServerWithPlainConfiguration(port: Int): Unit = runBlocking { install(SSE) routing { sse("/sse") { - val transport = SSEServerTransport("/message", this) + val transport = SseServerTransport("/message", this) val server = configureServer() // For SSE, you can also add prompts/tools/resources if needed: @@ -154,7 +148,7 @@ fun runSseMcpServerWithPlainConfiguration(port: Int): Unit = runBlocking { servers[transport.sessionId] = server - server.onCloseCallback = { + server.onClose { println("Server closed") servers.remove(transport.sessionId) } @@ -164,7 +158,7 @@ fun runSseMcpServerWithPlainConfiguration(port: Int): Unit = runBlocking { post("/message") { println("Received Message") val sessionId: String = call.request.queryParameters["sessionId"]!! - val transport = servers[sessionId]?.transport as? SSEServerTransport + val transport = servers[sessionId]?.transport as? SseServerTransport if (transport == null) { call.respond(HttpStatusCode.NotFound, "Session not found") return@post @@ -189,8 +183,8 @@ fun runSseMcpServerUsingKtorPlugin(port: Int): Unit = runBlocking { println("Use inspector to connect to the http://localhost:$port/sse") embeddedServer(CIO, host = "0.0.0.0", port = port) { - MCP { - return@MCP configureServer() + mcp { + return@mcp configureServer() } }.start(wait = true) } \ No newline at end of file diff --git a/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt b/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt index e3f56f34..7fadedef 100644 --- a/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt +++ b/samples/weather-stdio-server/src/main/kotlin/io/modelcontextprotocol/sample/server/McpWeatherServer.kt @@ -121,7 +121,7 @@ fun `run mcp server`() { runBlocking { server.connect(transport) val done = Job() - server.onCloseCallback = { + server.onClose { done.complete() } done.join() diff --git a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt index 8724cf02..e24419e5 100644 --- a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt +++ b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt @@ -30,7 +30,6 @@ public class ServerOptions( * * @param serverInfo Information about this server implementation (name, version). * @param options Configuration options for the server. - * @param onCloseCallback A callback invoked when the server connection closes. */ public open class Server( private val serverInfo: Implementation, @@ -127,7 +126,6 @@ public open class Server( /** * Called when the server connection is closing. - * Invokes [onCloseCallback] if set. */ override fun onClose() { logger.info { "Server connection closing" } diff --git a/src/jvmTest/kotlin/InMemoryTransport.kt b/src/jvmTest/kotlin/InMemoryTransport.kt index 987d166e..ed849e2b 100644 --- a/src/jvmTest/kotlin/InMemoryTransport.kt +++ b/src/jvmTest/kotlin/InMemoryTransport.kt @@ -1,6 +1,5 @@ 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. diff --git a/src/jvmTest/kotlin/client/InMemoryTransportTest.kt b/src/jvmTest/kotlin/client/InMemoryTransportTest.kt index 90acd1ec..bfe01f60 100644 --- a/src/jvmTest/kotlin/client/InMemoryTransportTest.kt +++ b/src/jvmTest/kotlin/client/InMemoryTransportTest.kt @@ -96,11 +96,9 @@ class InMemoryTransportTest { clientTransport.close() assertThrows { - runBlocking { - clientTransport.send( - InitializedNotification().toJSON() - ) - } + clientTransport.send( + InitializedNotification().toJSON() + ) } } }