From 772195b5e6b0a71ca1deff0cf86057855bc99b07 Mon Sep 17 00:00:00 2001
From: Sergey Ignatov <sergey.ignatov@jetbrains.com>
Date: Mon, 23 Dec 2024 09:04:06 +0100
Subject: [PATCH] Drop generics from Protocol class (#30)

The generic types in the Protocol class were not providing significant value and made
the code more complex. This change simplifies the API by removing these generic
parameters while maintaining the same functionality.
---
 .../kotlin/sdk/client/Client.kt               | 40 ++---------------
 .../kotlin/sdk/server/Server.kt               |  7 ++-
 .../kotlin/sdk/shared/Protocol.kt             | 43 ++++++-------------
 3 files changed, 18 insertions(+), 72 deletions(-)

diff --git a/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt b/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt
index 4eccbfa..2efc9eb 100644
--- a/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt
+++ b/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt
@@ -1,42 +1,7 @@
 package io.modelcontextprotocol.kotlin.sdk.client
 
-import io.modelcontextprotocol.kotlin.sdk.CallToolRequest
-import io.modelcontextprotocol.kotlin.sdk.CallToolResult
-import io.modelcontextprotocol.kotlin.sdk.CallToolResultBase
-import io.modelcontextprotocol.kotlin.sdk.ClientCapabilities
-import io.modelcontextprotocol.kotlin.sdk.ClientNotification
-import io.modelcontextprotocol.kotlin.sdk.ClientRequest
-import io.modelcontextprotocol.kotlin.sdk.ClientResult
-import io.modelcontextprotocol.kotlin.sdk.CompatibilityCallToolResult
-import io.modelcontextprotocol.kotlin.sdk.CompleteRequest
-import io.modelcontextprotocol.kotlin.sdk.CompleteResult
-import io.modelcontextprotocol.kotlin.sdk.EmptyRequestResult
-import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
-import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
-import io.modelcontextprotocol.kotlin.sdk.Implementation
-import io.modelcontextprotocol.kotlin.sdk.InitializeRequest
-import io.modelcontextprotocol.kotlin.sdk.InitializeResult
-import io.modelcontextprotocol.kotlin.sdk.InitializedNotification
-import io.modelcontextprotocol.kotlin.sdk.LATEST_PROTOCOL_VERSION
-import io.modelcontextprotocol.kotlin.sdk.ListPromptsRequest
-import io.modelcontextprotocol.kotlin.sdk.ListPromptsResult
-import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesRequest
-import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesResult
-import io.modelcontextprotocol.kotlin.sdk.ListResourcesRequest
-import io.modelcontextprotocol.kotlin.sdk.ListResourcesResult
-import io.modelcontextprotocol.kotlin.sdk.ListToolsRequest
-import io.modelcontextprotocol.kotlin.sdk.ListToolsResult
-import io.modelcontextprotocol.kotlin.sdk.LoggingLevel
+import io.modelcontextprotocol.kotlin.sdk.*
 import io.modelcontextprotocol.kotlin.sdk.LoggingMessageNotification.SetLevelRequest
-import io.modelcontextprotocol.kotlin.sdk.Method
-import io.modelcontextprotocol.kotlin.sdk.PingRequest
-import io.modelcontextprotocol.kotlin.sdk.ReadResourceRequest
-import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
-import io.modelcontextprotocol.kotlin.sdk.RootsListChangedNotification
-import io.modelcontextprotocol.kotlin.sdk.SUPPORTED_PROTOCOL_VERSIONS
-import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
-import io.modelcontextprotocol.kotlin.sdk.SubscribeRequest
-import io.modelcontextprotocol.kotlin.sdk.UnsubscribeRequest
 import io.modelcontextprotocol.kotlin.sdk.shared.Protocol
 import io.modelcontextprotocol.kotlin.sdk.shared.ProtocolOptions
 import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
@@ -45,6 +10,7 @@ import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.JsonNull
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
+
 /**
  * Options for configuring the MCP client.
  *
@@ -70,7 +36,7 @@ public class ClientOptions(
 public open class Client(
     private val clientInfo: Implementation,
     options: ClientOptions = ClientOptions(),
-) : Protocol<ClientRequest, ClientNotification, ClientResult>(options) {
+) : Protocol(options) {
 
     private var serverCapabilities: ServerCapabilities? = null
     private var serverVersion: Implementation? = null
diff --git a/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt b/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt
index c3f6dcf..75c1578 100644
--- a/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt
+++ b/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt
@@ -1,13 +1,12 @@
 package io.modelcontextprotocol.kotlin.sdk.server
 
 import io.github.oshai.kotlinlogging.KotlinLogging
-import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.serialization.json.JsonObject
 import io.modelcontextprotocol.kotlin.sdk.*
 import io.modelcontextprotocol.kotlin.sdk.shared.Protocol
 import io.modelcontextprotocol.kotlin.sdk.shared.ProtocolOptions
 import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.serialization.json.JsonObject
 
 private val logger = KotlinLogging.logger {}
 
@@ -37,7 +36,7 @@ public open class Server(
     private val serverInfo: Implementation,
     options: ServerOptions,
     public var onCloseCallback: (() -> Unit)? = null
-) : Protocol<ServerRequest, ServerNotification, ServerResult>(options) {
+) : Protocol(options) {
 
     /**
      * The client's reported capabilities after initialization.
diff --git a/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt b/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt
index 291b253..1e04fa0 100644
--- a/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt
+++ b/src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt
@@ -1,24 +1,7 @@
 package io.modelcontextprotocol.kotlin.sdk.shared
 
 import io.github.oshai.kotlinlogging.KotlinLogging
-import io.modelcontextprotocol.kotlin.sdk.CancelledNotification
-import io.modelcontextprotocol.kotlin.sdk.EmptyRequestResult
-import io.modelcontextprotocol.kotlin.sdk.ErrorCode
-import io.modelcontextprotocol.kotlin.sdk.JSONRPCError
-import io.modelcontextprotocol.kotlin.sdk.JSONRPCNotification
-import io.modelcontextprotocol.kotlin.sdk.JSONRPCRequest
-import io.modelcontextprotocol.kotlin.sdk.JSONRPCResponse
-import io.modelcontextprotocol.kotlin.sdk.McpError
-import io.modelcontextprotocol.kotlin.sdk.Method
-import io.modelcontextprotocol.kotlin.sdk.Notification
-import io.modelcontextprotocol.kotlin.sdk.PingRequest
-import io.modelcontextprotocol.kotlin.sdk.Progress
-import io.modelcontextprotocol.kotlin.sdk.ProgressNotification
-import io.modelcontextprotocol.kotlin.sdk.Request
-import io.modelcontextprotocol.kotlin.sdk.RequestId
-import io.modelcontextprotocol.kotlin.sdk.RequestResult
-import io.modelcontextprotocol.kotlin.sdk.fromJSON
-import io.modelcontextprotocol.kotlin.sdk.toJSON
+import io.modelcontextprotocol.kotlin.sdk.*
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.TimeoutCancellationException
@@ -107,14 +90,14 @@ internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) }
  * Implements MCP protocol framing on top of a pluggable transport, including
  * features like request/response linking, notifications, and progress.
  */
-public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notification, SendResultT : RequestResult>(
+public abstract class Protocol(
     @PublishedApi internal val options: ProtocolOptions?,
 ) {
     public var transport: Transport? = null
         private set
 
     @PublishedApi
-    internal val requestHandlers: MutableMap<String, suspend (request: JSONRPCRequest, extra: RequestHandlerExtra) -> SendResultT?> =
+    internal val requestHandlers: MutableMap<String, suspend (request: JSONRPCRequest, extra: RequestHandlerExtra) -> RequestResult?> =
         mutableMapOf()
     public val notificationHandlers: MutableMap<String, suspend (notification: JSONRPCNotification) -> Unit> =
         mutableMapOf()
@@ -143,7 +126,7 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
     /**
      * A handler to invoke for any request types that do not have their own handler installed.
      */
-    public var fallbackRequestHandler: (suspend (request: JSONRPCRequest, extra: RequestHandlerExtra) -> SendResultT?)? =
+    public var fallbackRequestHandler: (suspend (request: JSONRPCRequest, extra: RequestHandlerExtra) -> RequestResult?)? =
         null
 
     /**
@@ -158,8 +141,7 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
         }
 
         setRequestHandler<PingRequest>(Method.Defined.Ping) { request, _ ->
-            @Suppress("UNCHECKED_CAST")
-            EmptyRequestResult() as SendResultT
+            EmptyRequestResult()
         }
     }
 
@@ -344,12 +326,12 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
     public abstract fun assertRequestHandlerCapability(method: Method)
 
     /**
-     * Sends a request and wait for a response.
+     * Sends a request and waits for a response.
      *
      * Do not use this method to emit notifications! Use notification() instead.
      */
     public suspend fun <T : RequestResult> request(
-        request: SendRequestT,
+        request: Request,
         options: RequestOptions? = null,
     ): T {
         LOGGER.trace { "Sending request: ${request.method}" }
@@ -427,7 +409,7 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
     /**
      * Emits a notification, which is a one-way message that does not expect a response.
      */
-    public suspend fun notification(notification: SendNotificationT) {
+    public suspend fun notification(notification: Notification) {
         LOGGER.trace { "Sending notification: ${notification.method}" }
         val transport = this.transport ?: error("Not connected")
         assertNotificationCapability(notification.method)
@@ -446,7 +428,7 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
      */
     public inline fun <reified T : Request> setRequestHandler(
         method: Method,
-        noinline block: suspend (T, RequestHandlerExtra) -> SendResultT?,
+        noinline block: suspend (T, RequestHandlerExtra) -> RequestResult?,
     ) {
         setRequestHandler(typeOf<T>(), method, block)
     }
@@ -455,7 +437,7 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
     internal fun <T : Request> setRequestHandler(
         requestType: KType,
         method: Method,
-        block: suspend (T, RequestHandlerExtra) -> SendResultT?,
+        block: suspend (T, RequestHandlerExtra) -> RequestResult?,
     ) {
         assertRequestHandlerCapability(method)
 
@@ -470,8 +452,7 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
                 EmptyRequestResult()
             }
 
-            @Suppress("UNCHECKED_CAST")
-            response as SendResultT
+            response
         }
     }
 
@@ -500,4 +481,4 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
     public fun removeNotificationHandler(method: Method) {
         this.notificationHandlers.remove(method.value)
     }
-}
+}
\ No newline at end of file