Refactor onCloseCallback to onClose and update related usages across the codebase (#62)

This commit is contained in:
Pavel Gorgulov 2025-03-26 15:04:28 +01:00 • committed by GitHub
commit 1c5e54936d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 13 additions and 24 deletions

View file

@ -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<String>) {
}
fun configureServer(): Server {
val def = CompletableDeferred<Unit>()
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)
}

View file

@ -121,7 +121,7 @@ fun `run mcp server`() {
runBlocking {
server.connect(transport)
val done = Job()
server.onCloseCallback = {
server.onClose {
done.complete()
}
done.join()

View file

@ -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" }

View file

@ -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.

View file

@ -96,11 +96,9 @@ class InMemoryTransportTest {
clientTransport.close()
assertThrows<IllegalStateException> {
runBlocking {
clientTransport.send(
InitializedNotification().toJSON()
)
}
clientTransport.send(
InitializedNotification().toJSON()
)
}
}
}