Fix response result deserialization to use the original request method (#601)

kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt

@@ -45,6 +45,8 @@ import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.decodeFromJsonElement
+import kotlinx.serialization.json.encodeToJsonElement
 import kotlin.io.encoding.Base64
 import kotlin.io.encoding.ExperimentalEncodingApi
 import kotlin.math.pow
@@ -418,6 +420,24 @@ public class StreamableHttpClientTransport(
     private fun JsonObject.stringValue(key: String): String? =
         (get(key) as? JsonPrimitive)?.takeIf { it.isString }?.content
 
+    /**
+     * Decodes an incoming message, rewriting a response's id to [replayMessageId] when set.
+     *
+     * The id is rewritten in the raw JSON before decoding rather than via `JSONRPCResponse.copy`:
+     * decoding captures the raw result JSON, and a data-class copy would silently drop it, forcing
+     * Protocol to fall back to shape-based result decoding
+     * (see https://github.com/modelcontextprotocol/kotlin-sdk/issues/601).
+     */
+    private fun decodeIncomingMessage(json: String, replayMessageId: RequestId?): JSONRPCMessage {
+        val element = McpJson.parseToJsonElement(json)
+        val jsonObject = element as? JsonObject
+        if (replayMessageId == null || jsonObject == null || "result" !in jsonObject || "id" !in jsonObject) {
+            return McpJson.decodeFromJsonElement<JSONRPCMessage>(element)
+        }
+        val rewritten = JsonObject(jsonObject + ("id" to McpJson.encodeToJsonElement(replayMessageId)))
+        return McpJson.decodeFromJsonElement<JSONRPCMessage>(rewritten)
+    }
+
     @OptIn(ExperimentalEncodingApi::class)
     private fun String.encodeMcpHeaderValue(): String {
         val containsUnsafeCharacters = any { it != '\t' && it.code !in 0x20..0x7e }
@@ -450,14 +470,10 @@ public class StreamableHttpClientTransport(
                 when (event.event) {
                     null, "message" ->
                         event.data?.takeIf { it.isNotEmpty() }?.let { json ->
-                            runCatching { McpJson.decodeFromString<JSONRPCMessage>(json) }
+                            runCatching { decodeIncomingMessage(json, replayMessageId) }
                                 .onSuccess { msg ->
                                     if (msg is JSONRPCResponse) receivedResponse = true
-                                    if (replayMessageId != null && msg is JSONRPCResponse) {
-                                        _onMessage(msg.copy(id = replayMessageId))
-                                    } else {
-                                        _onMessage(msg)
-                                    }
+                                    _onMessage(msg)
                                 }
                                 .onFailure(_onError)
                         }
@@ -499,14 +515,10 @@ public class StreamableHttpClientTransport(
                 return
             }
             if (eventName == null || eventName == "message") {
-                runCatching { McpJson.decodeFromString<JSONRPCMessage>(data) }
+                runCatching { decodeIncomingMessage(data, replayMessageId) }
                     .onSuccess { msg ->
                         if (msg is JSONRPCResponse) receivedResponse = true
-                        if (replayMessageId != null && msg is JSONRPCResponse) {
-                            _onMessage(msg.copy(id = replayMessageId))
-                        } else {
-                            _onMessage(msg)
-                        }
+                        _onMessage(msg)
                     }
                     .onFailure {
                         _onError(it)

kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/streamable/http/StreamableHttpClientTransportReplayTest.kt

@@ -0,0 +1,93 @@
+package io.modelcontextprotocol.kotlin.sdk.client.streamable.http
+
+import io.kotest.matchers.nulls.shouldNotBeNull
+import io.kotest.matchers.shouldBe
+import io.ktor.client.HttpClient
+import io.ktor.client.engine.mock.MockEngine
+import io.ktor.client.engine.mock.respond
+import io.ktor.client.plugins.sse.SSE
+import io.ktor.http.ContentType
+import io.ktor.http.HttpHeaders
+import io.ktor.http.HttpMethod
+import io.ktor.http.HttpStatusCode
+import io.ktor.http.headersOf
+import io.ktor.utils.io.ByteReadChannel
+import io.modelcontextprotocol.kotlin.sdk.client.StreamableHttpClientTransport
+import io.modelcontextprotocol.kotlin.sdk.shared.Protocol
+import io.modelcontextprotocol.kotlin.sdk.types.GetTaskPayloadRequest
+import io.modelcontextprotocol.kotlin.sdk.types.GetTaskPayloadRequestParams
+import io.modelcontextprotocol.kotlin.sdk.types.GetTaskPayloadResult
+import io.modelcontextprotocol.kotlin.sdk.types.Method
+import kotlinx.coroutines.runBlocking
+import kotlinx.serialization.json.boolean
+import kotlinx.serialization.json.jsonPrimitive
+import kotlin.test.Test
+import kotlin.time.Duration.Companion.seconds
+
+/**
+ * Regression test for https://github.com/modelcontextprotocol/kotlin-sdk/issues/601 on the
+ * Streamable HTTP SSE replay path.
+ *
+ * Lives in `jvmTest` (not `commonTest`) because it drives a full `Protocol.request` round trip
+ * over a MockEngine: under `runTest`'s virtual clock the request's internal timeout fires while
+ * the real-dispatcher HTTP delivery is still in flight, so the test needs `runBlocking`.
+ */
+class StreamableHttpClientTransportReplayTest {
+
+    @Test
+    fun `should deserialize a replayed SSE response by the original request method`() = runBlocking {
+        // When a response arrives on the SSE stream with a different id, the transport rewrites it
+        // to the POSTed request's id. The rewrite must preserve the raw result JSON so that
+        // Protocol can decode the result according to the original request's method instead of
+        // guessing from the JSON shape.
+        val protocol = object : Protocol(null) {
+            override fun assertCapabilityForMethod(method: Method) = Unit
+            override fun assertNotificationCapability(method: Method) = Unit
+            override fun assertRequestHandlerCapability(method: Method) = Unit
+        }
+
+        val mockEngine = MockEngine { request ->
+            if (request.method == HttpMethod.Post) {
+                val sseContent = buildString {
+                    appendLine("id: ev-1")
+                    appendLine("event: message")
+                    appendLine(
+                        """data: {"jsonrpc":"2.0","id":"server-side-id",""" +
+                            """"result":{"content":[{"type":"text","text":"task output"}],"isError":false}}""",
+                    )
+                    appendLine()
+                }
+
+                respond(
+                    content = ByteReadChannel(sseContent),
+                    status = HttpStatusCode.OK,
+                    headers = headersOf(
+                        HttpHeaders.ContentType,
+                        ContentType.Text.EventStream.toString(),
+                    ),
+                )
+            } else {
+                respond("", HttpStatusCode.OK)
+            }
+        }
+        val httpClient = HttpClient(mockEngine) {
+            install(SSE) {
+                reconnectionTime = 1.seconds
+            }
+        }
+        val transport = StreamableHttpClientTransport(httpClient, url = "http://localhost:8080/mcp")
+
+        protocol.connect(transport)
+
+        val result = protocol.request<GetTaskPayloadResult>(
+            GetTaskPayloadRequest(GetTaskPayloadRequestParams(taskId = "task-42")),
+        )
+
+        // The tasks/result payload is CallToolResult-shaped; it must still surface as
+        // GetTaskPayloadResult with the raw payload preserved.
+        result["content"].shouldNotBeNull()
+        result["isError"]?.jsonPrimitive?.boolean shouldBe false
+
+        transport.close()
+    }
+}

kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt

@@ -22,6 +22,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.Request
 import io.modelcontextprotocol.kotlin.sdk.types.RequestId
 import io.modelcontextprotocol.kotlin.sdk.types.RequestResult
 import io.modelcontextprotocol.kotlin.sdk.types.fromJSON
+import io.modelcontextprotocol.kotlin.sdk.types.selectRequestResultDeserializer
 import io.modelcontextprotocol.kotlin.sdk.types.toJSON
 import kotlinx.atomicfu.AtomicRef
 import kotlinx.atomicfu.atomic
@@ -51,6 +52,7 @@ import kotlinx.coroutines.withTimeoutOrNull
 import kotlinx.serialization.SerializationException
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.decodeFromJsonElement
 import kotlinx.serialization.json.encodeToJsonElement
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.ContinuationInterceptor
@@ -739,6 +741,34 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
         }
     }
 
+    /**
+     * Decodes the result of an incoming [response] for the original request [method].
+     *
+     * The shape-decoded [JSONRPCResponse.result] is only a heuristic: distinct result types can
+     * share a JSON shape (e.g. a `tasks/result` payload vs. `CallToolResult`), so when the raw
+     * wire JSON was captured during transport decoding it is decoded again here with the result
+     * type declared by [method] (see
+     * https://github.com/modelcontextprotocol/kotlin-sdk/issues/601).
+     *
+     * Falls back to the shape-decoded [JSONRPCResponse.result] when no raw JSON is available
+     * (programmatically constructed responses), when the method is custom or unknown, or when the
+     * payload does not match the method's declared result type (e.g. a task-augmented `tools/call`
+     * returning `CreateTaskResult`).
+     */
+    private fun deserializeResult(method: Method, response: JSONRPCResponse): RequestResult {
+        val rawResult = response.rawResult ?: return response.result
+        val deserializer = selectRequestResultDeserializer(method.value) ?: return response.result
+        return try {
+            McpJson.decodeFromJsonElement(deserializer, rawResult)
+        } catch (e: SerializationException) {
+            logger.debug(e) {
+                "Failed to deserialize the result of '${method.value}' as " +
+                    "'${deserializer.descriptor.serialName}'; falling back to the shape-decoded result"
+            }
+            response.result
+        }
+    }
+
     /**
      * Closes the connection.
      */
@@ -813,7 +843,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
 
                 try {
                     @Suppress("UNCHECKED_CAST")
-                    result.complete(response!!.result as T)
+                    result.complete(deserializeResult(request.method, response!!) as T)
                 } catch (e: Throwable) {
                     result.completeExceptionally(e)
                 }

kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/jsonRpc.kt

@@ -5,6 +5,7 @@ package io.modelcontextprotocol.kotlin.sdk.types
 import kotlinx.serialization.EncodeDefault
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
+import kotlinx.serialization.Transient
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.decodeFromJsonElement
 import kotlinx.serialization.json.encodeToJsonElement
@@ -214,6 +215,18 @@ public data class JSONRPCResponse(val id: RequestId, val result: RequestResult =
     /** Always `"2.0"` to indicate JSON-RPC 2.0 protocol. */
     @EncodeDefault
     override val jsonrpc: String = JSONRPC_VERSION
+
+    /**
+     * The raw JSON of [result] exactly as received on the wire, captured when this response is
+     * decoded as part of a [JSONRPCMessage]. `null` for programmatically constructed responses.
+     *
+     * The shape-decoded [result] is only a heuristic: distinct result types can share a JSON
+     * shape (e.g. a `tasks/result` payload vs. [CallToolResult]). The request/response
+     * correlation point re-decodes [rawResult] with the result type declared by the original
+     * request's method (see https://github.com/modelcontextprotocol/kotlin-sdk/issues/601).
+     */
+    @Transient
+    internal var rawResult: JsonElement? = null
 }
 
 // ============================================================================

kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/serializers.kt

@@ -19,6 +19,7 @@ import kotlinx.serialization.json.JsonEncoder
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
 import kotlinx.serialization.json.contentOrNull
+import kotlinx.serialization.json.decodeFromJsonElement
 import kotlinx.serialization.json.jsonObject
 import kotlinx.serialization.json.jsonPrimitive
 import kotlinx.serialization.json.longOrNull
@@ -446,6 +447,47 @@ private fun selectServerResultDeserializer(element: JsonElement): Deserializatio
     }
 }
 
+/**
+ * Result deserializers keyed by the original request method.
+ *
+ * Shape-based selection ([selectClientResultDeserializer] / [selectServerResultDeserializer]) is only a
+ * heuristic: distinct result types can share a JSON shape (e.g. a `tasks/result` payload vs.
+ * [CallToolResult]), so the correlation point prefers the type declared by the method that was called.
+ */
+@OptIn(ExperimentalMcpApi::class)
+private val requestResultDeserializers: Map<String, DeserializationStrategy<out RequestResult>> by lazy {
+    mapOf(
+        Method.Defined.Initialize.value to InitializeResult.serializer(),
+        Method.Defined.ServerDiscover.value to DiscoverResult.serializer(),
+        Method.Defined.Ping.value to EmptyResult.serializer(),
+        Method.Defined.ResourcesList.value to ListResourcesResult.serializer(),
+        Method.Defined.ResourcesTemplatesList.value to ListResourceTemplatesResult.serializer(),
+        Method.Defined.ResourcesRead.value to ReadResourceResult.serializer(),
+        Method.Defined.ResourcesSubscribe.value to EmptyResult.serializer(),
+        Method.Defined.ResourcesUnsubscribe.value to EmptyResult.serializer(),
+        Method.Defined.PromptsList.value to ListPromptsResult.serializer(),
+        Method.Defined.PromptsGet.value to GetPromptResult.serializer(),
+        Method.Defined.ToolsList.value to ListToolsResult.serializer(),
+        Method.Defined.ToolsCall.value to CallToolResult.serializer(),
+        Method.Defined.LoggingSetLevel.value to EmptyResult.serializer(),
+        Method.Defined.SamplingCreateMessage.value to CreateMessageResult.serializer(),
+        Method.Defined.CompletionComplete.value to CompleteResult.serializer(),
+        Method.Defined.RootsList.value to ListRootsResult.serializer(),
+        Method.Defined.ElicitationCreate.value to ElicitResult.serializer(),
+        Method.Defined.TasksGet.value to GetTaskResult.serializer(),
+        Method.Defined.TasksResult.value to GetTaskPayloadResult.serializer(),
+        Method.Defined.TasksList.value to ListTasksResult.serializer(),
+        Method.Defined.TasksCancel.value to GetTaskResult.serializer(),
+    )
+}
+
+/**
+ * Selects the deserializer for the result type declared by the given request method.
+ * Returns null for custom or unknown methods, where shape-based decoding remains the fallback.
+ */
+internal fun selectRequestResultDeserializer(method: String): DeserializationStrategy<out RequestResult>? =
+    requestResultDeserializers[method]
+
 /**
  * Polymorphic serializer for [RequestResult] types.
  * Supports both client and server results.
@@ -488,6 +530,51 @@ internal object ServerResultPolymorphicSerializer :
 // JSON-RPC Serializers
 // ============================================================================
 
+/**
+ * Wire deserializer for [JSONRPCResponse] used when decoding a full [JSONRPCMessage].
+ *
+ * Captures the raw `result` JSON into [JSONRPCResponse.rawResult] so the request/response
+ * correlation point can decode it with the result type declared by the original request's method
+ * (shape-based decoding alone picks the wrong runtime type when distinct result types share a JSON
+ * shape, e.g. a `tasks/result` payload — see
+ * https://github.com/modelcontextprotocol/kotlin-sdk/issues/601).
+ *
+ * The public [JSONRPCResponse.result] keeps the shape-decoded value for direct consumers; a result
+ * whose shape matches no known type surfaces as [GetTaskPayloadResult] (the SDK's raw-payload result
+ * type) instead of failing the whole message decode, leaving the final interpretation to the
+ * correlation point.
+ */
+internal object JSONRPCResponseWireSerializer : KSerializer<JSONRPCResponse> {
+    override val descriptor: SerialDescriptor = JSONRPCResponse.serializer().descriptor
+
+    override fun serialize(encoder: Encoder, value: JSONRPCResponse) {
+        JSONRPCResponse.serializer().serialize(encoder, value)
+    }
+
+    override fun deserialize(decoder: Decoder): JSONRPCResponse {
+        val jsonDecoder = decoder as? JsonDecoder
+            ?: throw SerializationException("JSONRPCResponseWireSerializer requires a Json decoder")
+        val json = jsonDecoder.json
+        val jsonObject = jsonDecoder.decodeJsonElement().jsonObject
+
+        val idElement = jsonObject["id"]
+            ?: throw SerializationException("Missing required 'id' field in JSONRPCResponse")
+        val id = json.decodeFromJsonElement(RequestId.serializer(), idElement)
+
+        val resultElement = jsonObject["result"]
+        val result = when {
+            resultElement == null -> EmptyResult()
+
+            else -> try {
+                json.decodeFromJsonElement(RequestResultPolymorphicSerializer, resultElement)
+            } catch (e: SerializationException) {
+                (resultElement as? JsonObject)?.let(::GetTaskPayloadResult) ?: throw e
+            }
+        }
+        return JSONRPCResponse(id, result).also { it.rawResult = resultElement }
+    }
+}
+
 /**
  * Polymorphic serializer for [JSONRPCMessage] types.
  * Determines the message type based on the presence of specific fields:
@@ -503,7 +590,7 @@ internal object JSONRPCMessagePolymorphicSerializer :
         val jsonObj = element.jsonObject
         return when {
             "error" in jsonObj -> JSONRPCError.serializer()
-            "result" in jsonObj && "id" in jsonObj -> JSONRPCResponse.serializer()
+            "result" in jsonObj && "id" in jsonObj -> JSONRPCResponseWireSerializer
             "result" in jsonObj && jsonObj["result"]?.jsonObject?.isEmpty() == true -> JSONRPCEmptyMessage.serializer()
             "method" in jsonObj && "id" in jsonObj -> JSONRPCRequest.serializer()
             "method" in jsonObj -> JSONRPCNotification.serializer()