From 8426228bf8bc205b73f703fd1dbc109830be214b Mon Sep 17 00:00:00 2001 From: Leonid Stashevsky Date: Tue, 14 Jan 2025 16:24:55 +0100 Subject: [PATCH] Migrate to kotlin style callbacks. Close #14 (#36) --- README.md | 2 +- .../kotlin/sdk/client/SSEClientTransport.kt | 42 +++++++++++++++---- .../kotlin/sdk/server/Server.kt | 36 ++++++++++++---- .../kotlin/sdk/shared/Transport.kt | 6 +-- 4 files changed, 64 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index ee5d2f8..d4edfef 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ Add the dependency: ```kotlin dependencies { - implementation("io.modelcontextprotocol:kotlin-sdk:0.2.0") + implementation("io.modelcontextprotocol:kotlin-sdk:0.3.0") } ``` diff --git a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt index c4eba4a..5a7a4f2 100644 --- a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt +++ b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt @@ -33,9 +33,9 @@ public class SSEClientTransport( private var session: ClientSSESession by Delegates.notNull() private val endpoint = CompletableDeferred() - override var onClose: (() -> Unit)? = null - override var onError: ((Throwable) -> Unit)? = null - override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null + private var _onClose: (() -> Unit) = {} + private var _onError: ((Throwable) -> Unit) = {} + private var _onMessage: (suspend ((JSONRPCMessage) -> Unit)) = {} private var job: Job? = null @@ -67,7 +67,7 @@ public class SSEClientTransport( when (event.event) { "error" -> { val e = IllegalStateException("SSE error: ${event.data}") - onError?.invoke(e) + _onError(e) throw e } @@ -84,7 +84,7 @@ public class SSEClientTransport( endpoint.complete(maybeEndpoint.toString()) } catch (e: Exception) { - onError?.invoke(e) + _onError(e) close() error(e) } @@ -93,9 +93,9 @@ public class SSEClientTransport( else -> { try { val message = McpJson.decodeFromString(event.data ?: "") - onMessage?.invoke(message) + _onMessage(message) } catch (e: Exception) { - onError?.invoke(e) + _onError(e) } } } @@ -122,7 +122,7 @@ public class SSEClientTransport( error("Error POSTing to endpoint (HTTP ${response.status}): $text") } } catch (e: Exception) { - onError?.invoke(e) + _onError(e) throw e } } @@ -133,7 +133,31 @@ public class SSEClientTransport( } session.cancel() - onClose?.invoke() + _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) + } + } } 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 1306f9f..8724cf0 100644 --- a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt +++ b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt @@ -35,8 +35,9 @@ public class ServerOptions( public open class Server( private val serverInfo: Implementation, options: ServerOptions, - public var onCloseCallback: (() -> Unit)? = null ) : Protocol(options) { + private var _onInitialized: (() -> Unit) = {} + private var _onClose: () -> Unit = {} /** * The client's reported capabilities after initialization. @@ -49,18 +50,13 @@ public open class Server( */ public var clientVersion: Implementation? = null private set + private val capabilities: ServerCapabilities = options.capabilities private val tools = mutableMapOf() private val prompts = mutableMapOf() private val resources = mutableMapOf() - /** - * A callback invoked when the server has completed the initialization sequence. - * After initialization, the server is ready to handle requests. - */ - public var onInitialized: (() -> Unit)? = null - init { logger.debug { "Initializing MCP server with capabilities: $capabilities" } @@ -69,7 +65,7 @@ public open class Server( handleInitialize(request) } setNotificationHandler(Method.Defined.NotificationsInitialized) { - onInitialized?.invoke() + _onInitialized() CompletableDeferred(Unit) } @@ -107,13 +103,35 @@ public open class Server( } } + /** + * Registers a callback to be invoked when the server has completed initialization. + */ + public fun onInitalized(block: () -> Unit) { + val old = _onInitialized + _onInitialized = { + old() + block() + } + } + + /** + * Registers a callback to be invoked when the server connection is closing. + */ + public fun onClose(block: () -> Unit) { + val old = _onClose + _onClose = { + old() + block() + } + } + /** * Called when the server connection is closing. * Invokes [onCloseCallback] if set. */ override fun onClose() { logger.info { "Server connection closing" } - onCloseCallback?.invoke() + _onClose() } /** diff --git a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt index 797a6cb..b1b15a4 100644 --- a/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt +++ b/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt @@ -31,7 +31,7 @@ public interface Transport { * * This should be invoked when close() is called as well. */ - public var onClose: (() -> Unit)? + public fun onClose(block: () -> Unit) /** * Callback for when an error occurs. @@ -39,10 +39,10 @@ public interface Transport { * Note that errors are not necessarily fatal; they are used for reporting any kind of * exceptional condition out of band. */ - public var onError: ((Throwable) -> Unit)? + public fun onError(block: (Throwable) -> Unit) /** * Callback for when a message (request or response) is received over the connection. */ - public var onMessage: (suspend ((JSONRPCMessage) -> Unit))? + public fun onMessage(block: suspend (JSONRPCMessage) -> Unit) }