Cleanup API (#29)

This commit is contained in:
Sergey Ignatov 2024-12-23 17:14:24 +01:00 • committed by GitHub
commit 67d1f7775a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 18 additions and 22 deletions

View file

@ -111,7 +111,7 @@ public open class Server(
* Called when the server connection is closing. * Called when the server connection is closing.
* Invokes [onCloseCallback] if set. * Invokes [onCloseCallback] if set.
*/ */
override fun onclose() { override fun onClose() {
logger.info { "Server connection closing" } logger.info { "Server connection closing" }
onCloseCallback?.invoke() onCloseCallback?.invoke()
} }

View file

@ -47,7 +47,7 @@ public open class ProtocolOptions(
* Note that this DOES NOT affect checking of _local_ side capabilities, as it is * Note that this DOES NOT affect checking of _local_ side capabilities, as it is
* considered a logic error to mis-specify those. * considered a logic error to mis-specify those.
* *
* Currently this defaults to false, for backwards compatibility with SDK versions * Currently, this defaults to false, for backwards compatibility with SDK versions
* that did not advertise capabilities correctly. In future, this will default to true. * that did not advertise capabilities correctly. In future, this will default to true.
*/ */
public var enforceStrictCapabilities: Boolean = false, public var enforceStrictCapabilities: Boolean = false,
@ -114,14 +114,14 @@ public abstract class Protocol(
* *
* This is invoked when close() is called as well. * This is invoked when close() is called as well.
*/ */
public open fun onclose() {} public open fun onClose() {}
/** /**
* Callback for when an error occurs. * Callback for when an error occurs.
* *
* Note that errors are not necessarily fatal they are used for reporting any kind of exceptional condition out of band. * Note that errors are not necessarily fatal they are used for reporting any kind of exceptional condition out of band.
*/ */
public open fun onerror(error: Throwable) {} public open fun onError(error: Throwable) {}
/** /**
* A handler to invoke for any request types that do not have their own handler installed. * A handler to invoke for any request types that do not have their own handler installed.
@ -136,7 +136,7 @@ public abstract class Protocol(
init { init {
setNotificationHandler<ProgressNotification>(Method.Defined.NotificationsProgress) { notification -> setNotificationHandler<ProgressNotification>(Method.Defined.NotificationsProgress) { notification ->
this.onProgress(notification) onProgress(notification)
COMPLETED COMPLETED
} }
@ -153,11 +153,11 @@ public abstract class Protocol(
public open suspend fun connect(transport: Transport) { public open suspend fun connect(transport: Transport) {
this.transport = transport this.transport = transport
transport.onClose = { transport.onClose = {
this.onClose() doClose()
} }
transport.onError = { transport.onError = {
this.onError(it) onError(it)
} }
transport.onMessage = { message -> transport.onMessage = { message ->
@ -172,11 +172,11 @@ public abstract class Protocol(
return transport.start() return transport.start()
} }
private fun onClose() { private fun doClose() {
responseHandlers.clear() responseHandlers.clear()
progressHandlers.clear() progressHandlers.clear()
transport = null transport = null
onclose() onClose()
val error = McpError(ErrorCode.Defined.ConnectionClosed.code, "Connection closed") val error = McpError(ErrorCode.Defined.ConnectionClosed.code, "Connection closed")
for (handler in responseHandlers.values) { for (handler in responseHandlers.values) {
@ -184,10 +184,6 @@ public abstract class Protocol(
} }
} }
private fun onError(error: Throwable) {
onerror(error)
}
private suspend fun onNotification(notification: JSONRPCNotification) { private suspend fun onNotification(notification: JSONRPCNotification) {
LOGGER.trace { "Received notification: ${notification.method}" } LOGGER.trace { "Received notification: ${notification.method}" }
val function = notificationHandlers[notification.method] val function = notificationHandlers[notification.method]
@ -208,7 +204,7 @@ public abstract class Protocol(
private suspend fun onRequest(request: JSONRPCRequest) { private suspend fun onRequest(request: JSONRPCRequest) {
LOGGER.trace { "Received request: ${request.method} (id: ${request.id})" } LOGGER.trace { "Received request: ${request.method} (id: ${request.id})" }
val handler = requestHandlers[request.method] ?: this.fallbackRequestHandler val handler = requestHandlers[request.method] ?: fallbackRequestHandler
if (handler === null) { if (handler === null) {
LOGGER.trace { "No handler found for request: ${request.method}" } LOGGER.trace { "No handler found for request: ${request.method}" }
@ -260,13 +256,13 @@ public abstract class Protocol(
val total = notification.total val total = notification.total
val progressToken = notification.progressToken val progressToken = notification.progressToken
val handler = this.progressHandlers[progressToken] val handler = progressHandlers[progressToken]
if (handler == null) { if (handler == null) {
val error = Error( val error = Error(
"Received a progress notification for an unknown token: ${McpJson.encodeToString(notification)}", "Received a progress notification for an unknown token: ${McpJson.encodeToString(notification)}",
) )
LOGGER.error { error.message } LOGGER.error { error.message }
this.onError(error) onError(error)
return return
} }
@ -275,14 +271,14 @@ public abstract class Protocol(
private fun onResponse(response: JSONRPCResponse?, error: JSONRPCError?) { private fun onResponse(response: JSONRPCResponse?, error: JSONRPCError?) {
val messageId = response?.id val messageId = response?.id
val handler = this.responseHandlers[messageId] val handler = responseHandlers[messageId]
if (handler == null) { if (handler == null) {
this.onError(Error("Received a response for an unknown message ID: ${McpJson.encodeToString(response)}")) onError(Error("Received a response for an unknown message ID: ${McpJson.encodeToString(response)}"))
return return
} }
this.responseHandlers.remove(messageId) responseHandlers.remove(messageId)
this.progressHandlers.remove(messageId) progressHandlers.remove(messageId)
if (response != null) { if (response != null) {
handler(response, null) handler(response, null)
} else { } else {
@ -469,7 +465,7 @@ public abstract class Protocol(
* Note that this will replace any previous notification handler for the same method. * Note that this will replace any previous notification handler for the same method.
*/ */
public fun <T : Notification> setNotificationHandler(method: Method, handler: (notification: T) -> Deferred<Unit>) { public fun <T : Notification> setNotificationHandler(method: Method, handler: (notification: T) -> Deferred<Unit>) {
this.notificationHandlers[method.value] = { notificationHandlers[method.value] = {
@Suppress("UNCHECKED_CAST") @Suppress("UNCHECKED_CAST")
handler(it.fromJSON() as T) handler(it.fromJSON() as T)
} }
@ -479,6 +475,6 @@ public abstract class Protocol(
* Removes the notification handler for the given method. * Removes the notification handler for the given method.
*/ */
public fun removeNotificationHandler(method: Method) { public fun removeNotificationHandler(method: Method) {
this.notificationHandlers.remove(method.value) notificationHandlers.remove(method.value)
} }
} }