Accept requestID as a string #25 (#26)

Author: ignatov

src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt

@@ -15,6 +15,7 @@ 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
@@ -119,11 +120,11 @@ public abstract class Protocol<SendRequestT : Request, SendNotificationT : Notif
         mutableMapOf()
 
     @PublishedApi
-    internal val responseHandlers: MutableMap<Long, (response: JSONRPCResponse?, error: Exception?) -> Unit> =
+    internal val responseHandlers: MutableMap<RequestId, (response: JSONRPCResponse?, error: Exception?) -> Unit> =
         mutableMapOf()
 
     @PublishedApi
-    internal val progressHandlers: MutableMap<Long, ProgressCallback> = mutableMapOf()
+    internal val progressHandlers: MutableMap<RequestId, ProgressCallback> = mutableMapOf()
 
     /**
      * Callback for when the connection is closed for any reason.

src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/types.kt

@@ -25,7 +25,7 @@ private val REQUEST_MESSAGE_ID = AtomicLong(0L)
  * A progress token, used to associate progress notifications with the original request.
  * Stores message ID.
  */
-public typealias ProgressToken = Long
+public typealias ProgressToken = RequestId
 
 /**
  * An opaque token used to represent a cursor for pagination.
@@ -191,7 +191,14 @@ public data class EmptyRequestResult(
 /**
  * A uniquely identifying ID for a request in JSON-RPC.
  */
-public typealias RequestId = Long
+@Serializable(with = RequestIdSerializer::class)
+public sealed interface RequestId {
+    @Serializable
+    public data class StringId(val value: String) : RequestId
+
+    @Serializable
+    public data class NumberId(val value: Long) : RequestId
+}
 
 /**
  * Represents a JSON-RPC message in the protocol.
@@ -204,7 +211,7 @@ public sealed interface JSONRPCMessage
  */
 @Serializable
 public data class JSONRPCRequest(
-    val id: RequestId = REQUEST_MESSAGE_ID.incrementAndGet(),
+    val id: RequestId = RequestId.NumberId(REQUEST_MESSAGE_ID.incrementAndGet()),
     val method: String,
     val params: JsonElement? = null,
     val jsonrpc: String = JSONRPC_VERSION,

src/main/kotlin/io/modelcontextprotocol/kotlin/sdk/types.util.kt

@@ -7,6 +7,7 @@ import kotlinx.serialization.KSerializer
 import kotlinx.serialization.descriptors.PrimitiveKind
 import kotlinx.serialization.descriptors.PrimitiveSerialDescriptor
 import kotlinx.serialization.descriptors.SerialDescriptor
+import kotlinx.serialization.descriptors.buildClassSerialDescriptor
 import kotlinx.serialization.encoding.Decoder
 import kotlinx.serialization.encoding.Encoder
 import kotlinx.serialization.json.*
@@ -270,3 +271,29 @@ internal object JSONRPCMessagePolymorphicSerializer :
 }
 
 internal val EmptyJsonObject = JsonObject(emptyMap())
+
+public class RequestIdSerializer : KSerializer<RequestId> {
+    override val descriptor: SerialDescriptor = buildClassSerialDescriptor("RequestId")
+
+    override fun deserialize(decoder: Decoder): RequestId {
+        val jsonDecoder = decoder as? JsonDecoder ?: error("Can only deserialize JSON")
+        val element = jsonDecoder.decodeJsonElement()
+
+        return when (element) {
+            is JsonPrimitive -> when {
+                element.isString -> RequestId.StringId(element.content)
+                element.longOrNull != null -> RequestId.NumberId(element.long)
+                else -> error("Invalid RequestId type")
+            }
+            else -> error("Invalid RequestId format")
+        }
+    }
+
+    override fun serialize(encoder: Encoder, value: RequestId) {
+        val jsonEncoder = encoder as? JsonEncoder ?: error("Can only serialize JSON")
+        when (value) {
+            is RequestId.StringId -> jsonEncoder.encodeString(value.value)
+            is RequestId.NumberId -> jsonEncoder.encodeLong(value.value)
+        }
+    }
+}

src/test/kotlin/client/TypesTest.kt

@@ -24,4 +24,24 @@ class TypesTest {
         val line = "{\"result\":{\"content\":[{\"type\":\"text\"}],\"isError\":false},\"jsonrpc\":\"2.0\",\"id\":4}"
         McpJson.decodeFromString<JSONRPCMessage>(line)
     }
+
+    @Test
+    fun testJSONRPCMessageWithStringId() {
+        val line = """
+            {
+              "jsonrpc": "2.0",
+              "method": "initialize",
+              "id": "ebf9f64a-0",
+              "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {
+                  "name": "mcp-java-client",
+                  "version": "0.2.0"
+                }
+              }
+            }
+        """.trimIndent()
+        McpJson.decodeFromString<JSONRPCMessage>(line)
+    }
 }