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

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

@@ -3,6 +3,7 @@ package io.modelcontextprotocol.kotlin.sdk.shared
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotification
 import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams
+import io.modelcontextprotocol.kotlin.sdk.types.CreateTaskResult
 import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult
 import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCEmptyMessage
 import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError
@@ -21,7 +22,9 @@ import io.modelcontextprotocol.kotlin.sdk.types.RPCError
 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.RequestResultPolymorphicSerializer
 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 +54,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 +743,34 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
         }
     }
 
+    /**
+     * Decodes the result of an incoming [response] for the original [request].
+     *
+     * 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 its 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). Custom or unknown methods use shape-based decoding,
+     * including its validation errors, instead of accepting the wire decoder's raw-payload fallback. A request
+     * with task augmentation may receive a [CreateTaskResult] envelope instead of the method's normal
+     * result. Decode that envelope directly from the raw JSON before considering the normal result,
+     * since its extension fields can overlap with other result shapes.
+     * All other failures to decode the declared result type are propagated to the caller.
+     */
+    private fun deserializeResult(request: JSONRPCRequest, response: JSONRPCResponse): RequestResult {
+        val rawResult = response.rawResult ?: return response.result
+        val deserializer = selectRequestResultDeserializer(request.method)
+            ?: return McpJson.decodeFromJsonElement(RequestResultPolymorphicSerializer, rawResult)
+        val taskRequested = (request.params as? JsonObject)?.get("task") is JsonObject
+        if (taskRequested && rawResult is JsonObject && "task" in rawResult) {
+            return McpJson.decodeFromJsonElement(CreateTaskResult.serializer(), rawResult)
+        }
+        return McpJson.decodeFromJsonElement(deserializer, rawResult)
+    }
+
     /**
      * Closes the connection.
      */
@@ -813,7 +845,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
 
                 try {
                     @Suppress("UNCHECKED_CAST")
-                    result.complete(response!!.result as T)
+                    result.complete(deserializeResult(jsonRpcRequest, 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
@@ -210,10 +211,44 @@ public data class JSONRPCNotification(val method: String, val params: JsonElemen
  * @property result The result of the method execution. Structure depends on the method that was called.
  */
 @Serializable
-public data class JSONRPCResponse(val id: RequestId, val result: RequestResult = EmptyResult()) : JSONRPCMessage {
+@ConsistentCopyVisibility
+public data class JSONRPCResponse private constructor(
+    public val id: RequestId,
+    public val result: RequestResult = EmptyResult(),
+    /** Raw wire JSON for method-based decoding; excluded from the public response value. */
+    @Transient
+    internal var rawResult: JsonElement? = null,
+) : JSONRPCMessage {
+    /**
+     * Creates a response with the original public constructor interface.
+     *
+     * @param id The identifier of the original request.
+     * @param result The result of the method execution.
+     */
+    public constructor(id: RequestId, result: RequestResult = EmptyResult()) : this(id, result, null)
+
     /** Always `"2.0"` to indicate JSON-RPC 2.0 protocol. */
     @EncodeDefault
     override val jsonrpc: String = JSONRPC_VERSION
+
+    /**
+     * Copies this response, preserving the original wire JSON when [result] is unchanged.
+     * Replacing the result discards the old wire JSON so it cannot override the new value.
+     *
+     * @param id The identifier for the copied response.
+     * @param result The result for the copied response.
+     * @return A response carrying the selected public fields and matching wire data.
+     */
+    public fun copy(id: RequestId = this.id, result: RequestResult = this.result): JSONRPCResponse =
+        JSONRPCResponse(id, result, rawResult.takeIf { result === this.result })
+
+    // Keep wire data out of the value semantics of the original two-field data class.
+    override fun equals(other: Any?): Boolean =
+        this === other || (other is JSONRPCResponse && id == other.id && result == other.result)
+
+    override fun hashCode(): Int = 31 * id.hashCode() + result.hashCode()
+
+    override fun toString(): String = "JSONRPCResponse(id=$id, result=$result)"
 }
 
 // ============================================================================

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

@@ -21,6 +21,7 @@ import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.JsonPrimitive
 import kotlinx.serialization.json.JsonTransformingSerializer
 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
@@ -471,6 +472,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.
@@ -513,6 +555,52 @@ 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 when
+ * decoding succeeds. Object results that fail these provisional shape checks are carried as
+ * [GetTaskPayloadResult] (the SDK's raw-payload result type), leaving the final interpretation and
+ * validation 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: IllegalArgumentException) {
+                // Shape checks can fail before serializer validation, e.g. content without a type.
+                (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:
@@ -528,7 +616,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()