From a01b925c3a2a7471a43e597c74584c41d43c2929 Mon Sep 17 00:00:00 2001 From: King Star Date: Thu, 23 Jul 2026 12:56:21 +0800 Subject: [PATCH] fix(core): disambiguate task result responses (#601) --- .../kotlin/sdk/shared/Protocol.kt | 20 ++- .../kotlin/sdk/types/jsonRpc.kt | 10 ++ .../kotlin/sdk/types/serializers.kt | 72 ++++++++-- .../kotlin/sdk/shared/ProtocolTest.kt | 133 ++++++++++++++++++ .../kotlin/sdk/types/JsonRpcTest.kt | 19 +++ 5 files changed, 241 insertions(+), 13 deletions(-) diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt index 908e524ce..af573a20e 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt @@ -4,6 +4,8 @@ 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.EmptyResult +import io.modelcontextprotocol.kotlin.sdk.types.GetTaskPayloadResult +import io.modelcontextprotocol.kotlin.sdk.types.GetTaskResult import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCEmptyMessage import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification @@ -51,7 +53,9 @@ 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 kotlinx.serialization.json.jsonObject import kotlin.coroutines.AbstractCoroutineContextElement import kotlin.coroutines.ContinuationInterceptor import kotlin.coroutines.CoroutineContext @@ -813,7 +817,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio try { @Suppress("UNCHECKED_CAST") - result.complete(response!!.result as T) + result.complete(response!!.resultForRequest(request) as T) } catch (e: Throwable) { result.completeExceptionally(e) } @@ -901,6 +905,20 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } } + private fun JSONRPCResponse.resultForRequest(request: Request): RequestResult { + val rawResult = rawResultOrNull() ?: McpJson.encodeToJsonElement(result) + return when (request.method.value) { + Method.Defined.TasksGet.value, + Method.Defined.TasksCancel.value, + -> McpJson.decodeFromJsonElement(rawResult) + + Method.Defined.TasksResult.value, + -> GetTaskPayloadResult(rawResult.jsonObject) + + else -> result + } + } + /** * Emits a notification, which is a one-way message that does not expect a response. */ diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/jsonRpc.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/jsonRpc.kt index 7e10904a9..5e74d73d4 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/jsonRpc.kt +++ b/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,15 @@ 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 + + @Transient + private var rawResult: JsonElement? = null + + internal fun withRawResult(result: JsonElement): JSONRPCResponse = apply { + rawResult = result + } + + internal fun rawResultOrNull(): JsonElement? = rawResult } // ============================================================================ diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/serializers.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/serializers.kt index 36b7e372d..a0cac5f0a 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/serializers.kt +++ b/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 @@ -454,12 +455,15 @@ private fun selectServerResultDeserializer(element: JsonElement): Deserializatio internal object RequestResultPolymorphicSerializer : JsonContentPolymorphicSerializer(RequestResult::class) { override fun selectDeserializer(element: JsonElement): DeserializationStrategy = - selectClientResultDeserializer(element) - ?: selectServerResultDeserializer(element) - ?: selectEmptyResult(element) + selectRequestResultDeserializer(element) ?: throw SerializationException("Cannot determine RequestResult type from JSON: ${element.jsonObject.keys}") } +private fun selectRequestResultDeserializer(element: JsonElement): DeserializationStrategy? = + selectClientResultDeserializer(element) + ?: selectServerResultDeserializer(element) + ?: selectEmptyResult(element) + /** * Polymorphic serializer for [ClientResult] types. * Throws [SerializationException] if the result type cannot be determined. @@ -497,20 +501,64 @@ internal object ServerResultPolymorphicSerializer : * - "method" + "id" -> JSONRPCRequest * - "method" -> JSONRPCNotification */ -internal object JSONRPCMessagePolymorphicSerializer : - JsonContentPolymorphicSerializer(JSONRPCMessage::class) { - override fun selectDeserializer(element: JsonElement): DeserializationStrategy { +internal object JSONRPCMessagePolymorphicSerializer : KSerializer { + override val descriptor: SerialDescriptor = JSONRPCResponse.serializer().descriptor + + override fun serialize(encoder: Encoder, value: JSONRPCMessage) { + when (value) { + is JSONRPCResponse -> encoder.encodeSerializableValue(JSONRPCResponse.serializer(), value) + + is JSONRPCError -> encoder.encodeSerializableValue(JSONRPCError.serializer(), value) + + is JSONRPCRequest -> encoder.encodeSerializableValue(JSONRPCRequest.serializer(), value) + + is JSONRPCNotification -> encoder.encodeSerializableValue(JSONRPCNotification.serializer(), value) + + JSONRPCEmptyMessage -> encoder.encodeSerializableValue( + JSONRPCEmptyMessage.serializer(), + JSONRPCEmptyMessage, + ) + } + } + + override fun deserialize(decoder: Decoder): JSONRPCMessage { + val jsonDecoder = decoder as? JsonDecoder + ?: throw SerializationException("JSONRPCMessagePolymorphicSerializer requires a Json decoder") + val element = jsonDecoder.decodeJsonElement() val jsonObj = element.jsonObject return when { - "error" in jsonObj -> JSONRPCError.serializer() - "result" in jsonObj && "id" in jsonObj -> JSONRPCResponse.serializer() - "result" in jsonObj && jsonObj["result"]?.jsonObject?.isEmpty() == true -> JSONRPCEmptyMessage.serializer() - "method" in jsonObj && "id" in jsonObj -> JSONRPCRequest.serializer() - "method" in jsonObj -> JSONRPCNotification.serializer() - jsonObj.isEmpty() || jsonObj.keys == setOf("jsonrpc") -> JSONRPCEmptyMessage.serializer() + "error" in jsonObj -> jsonDecoder.json.decodeFromJsonElement(JSONRPCError.serializer(), element) + + "result" in jsonObj && "id" in jsonObj -> decodeResponse(jsonDecoder, element) + + "result" in jsonObj && jsonObj["result"]?.jsonObject?.isEmpty() == true -> + jsonDecoder.json.decodeFromJsonElement(JSONRPCEmptyMessage.serializer(), element) + + "method" in jsonObj && "id" in jsonObj -> + jsonDecoder.json.decodeFromJsonElement(JSONRPCRequest.serializer(), element) + + "method" in jsonObj -> jsonDecoder.json.decodeFromJsonElement(JSONRPCNotification.serializer(), element) + + jsonObj.isEmpty() || jsonObj.keys == setOf("jsonrpc") -> + jsonDecoder.json.decodeFromJsonElement(JSONRPCEmptyMessage.serializer(), element) + else -> throw SerializationException("Invalid JSONRPCMessage type: ${jsonObj.keys}") } } + + private fun decodeResponse(jsonDecoder: JsonDecoder, element: JsonElement): JSONRPCResponse { + val jsonObject = element.jsonObject + val rawResult = jsonObject.getValue("result") + val response = if (rawResult is JsonObject && selectRequestResultDeserializer(rawResult) == null) { + JSONRPCResponse( + id = jsonDecoder.json.decodeFromJsonElement(jsonObject.getValue("id")), + result = GetTaskPayloadResult(rawResult), + ) + } else { + jsonDecoder.json.decodeFromJsonElement(JSONRPCResponse.serializer(), element) + } + return response.withRawResult(rawResult) + } } /** diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt index 2846cccaa..12930134d 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTest.kt @@ -8,8 +8,17 @@ import io.kotest.matchers.shouldBe import io.kotest.matchers.shouldNotBe import io.kotest.matchers.string.shouldContain import io.kotest.matchers.types.shouldBeInstanceOf +import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult +import io.modelcontextprotocol.kotlin.sdk.types.CancelTaskRequest +import io.modelcontextprotocol.kotlin.sdk.types.CancelTaskRequestParams import io.modelcontextprotocol.kotlin.sdk.types.CustomRequest import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult +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.GetTaskRequest +import io.modelcontextprotocol.kotlin.sdk.types.GetTaskRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.GetTaskResult import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest @@ -24,13 +33,17 @@ import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequest import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequestParams import io.modelcontextprotocol.kotlin.sdk.types.RequestId import io.modelcontextprotocol.kotlin.sdk.types.RequestMeta +import io.modelcontextprotocol.kotlin.sdk.types.TaskStatus import kotlinx.coroutines.async import kotlinx.coroutines.launch import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.JsonArray +import kotlinx.serialization.json.JsonElement import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonObjectBuilder import kotlinx.serialization.json.JsonPrimitive import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.decodeFromJsonElement import kotlinx.serialization.json.encodeToJsonElement import kotlinx.serialization.json.int import kotlinx.serialization.json.jsonObject @@ -195,6 +208,110 @@ class ProtocolTest { inFlight.await() } + @Test + fun `tasks result preserves payloads that overlap known result shapes`() = runTest { + protocol.connect(transport) + val payload = buildJsonObject { + put("content", JsonArray(emptyList())) + put("isError", JsonPrimitive(false)) + put("extension", JsonPrimitive("preserved")) + } + + val inFlight = async { + protocol.request( + CustomRequest(method = Method.Custom(Method.Defined.TasksResult.value), params = null), + ) + } + + val sent = transport.awaitRequest() + transport.deliver(decodeResponse(sent.id, payload)) + + inFlight.await().json shouldBe payload + } + + @Test + fun `tasks result preserves payloads that resemble task state`() = runTest { + protocol.connect(transport) + val payload = McpJson.encodeToJsonElement(completedTask("nested-task")).jsonObject + + val inFlight = async { + protocol.request( + GetTaskPayloadRequest(GetTaskPayloadRequestParams("task-1")), + ) + } + + val sent = transport.awaitRequest() + transport.deliver(decodeResponse(sent.id, payload)) + + inFlight.await().json shouldBe payload + } + + @Test + fun `tasks result preserves otherwise unknown payloads`() = runTest { + protocol.connect(transport) + val payload = buildJsonObject { + put("extension", JsonPrimitive("preserved")) + } + + val inFlight = async { + protocol.request( + GetTaskPayloadRequest(GetTaskPayloadRequestParams("task-1")), + ) + } + + val sent = transport.awaitRequest() + transport.deliver(decodeResponse(sent.id, payload)) + + inFlight.await().json shouldBe payload + } + + @Test + fun `tasks result adapts directly constructed transport responses`() = runTest { + protocol.connect(transport) + val responseResult = CallToolResult(content = emptyList(), isError = false) + + val inFlight = async { + protocol.request( + GetTaskPayloadRequest(GetTaskPayloadRequestParams("task-1")), + ) + } + + val sent = transport.awaitRequest() + transport.deliver(JSONRPCResponse(sent.id, responseResult)) + + inFlight.await().json shouldBe McpJson.encodeToJsonElement(responseResult).jsonObject + } + + @Test + fun `tasks get preserves typed task state responses`() = runTest { + protocol.connect(transport) + val expected = completedTask("task-1") + + val inFlight = async { + protocol.request(GetTaskRequest(GetTaskRequestParams("task-1"))) + } + + val sent = transport.awaitRequest() + transport.deliver(decodeResponse(sent.id, McpJson.encodeToJsonElement(expected))) + + inFlight.await() shouldBe expected + } + + @Test + fun `tasks cancel preserves typed task state responses`() = runTest { + protocol.connect(transport) + val expected = completedTask("task-1").copy(status = TaskStatus.Cancelled) + + val inFlight = async { + protocol.request(CancelTaskRequest(CancelTaskRequestParams("task-1"))) + } + + val sent = transport.awaitRequest() + transport.deliver(decodeResponse(sent.id, McpJson.encodeToJsonElement(expected))) + + inFlight.await() shouldBe expected + } + @Test fun `request handler receives enriched extra and ambient context element`() = runTest { protocol.connect(transport) @@ -368,3 +485,19 @@ class ProtocolTest { private fun metaOf(builderAction: JsonObjectBuilder.() -> Unit): RequestMeta = RequestMeta(metaJson(builderAction)) private fun metaJson(builderAction: JsonObjectBuilder.() -> Unit): JsonObject = buildJsonObject(builderAction) + +private fun decodeResponse(id: RequestId, result: JsonElement): JSONRPCMessage = McpJson.decodeFromJsonElement( + buildJsonObject { + put("jsonrpc", JsonPrimitive("2.0")) + put("id", McpJson.encodeToJsonElement(id)) + put("result", result) + }, +) + +private fun completedTask(taskId: String): GetTaskResult = GetTaskResult( + taskId = taskId, + status = TaskStatus.Completed, + createdAt = "2026-07-23T00:00:00Z", + lastUpdatedAt = "2026-07-23T00:00:01Z", + ttl = null, +) diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/JsonRpcTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/JsonRpcTest.kt index 7de360d4b..912e1d597 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/JsonRpcTest.kt +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/JsonRpcTest.kt @@ -435,6 +435,25 @@ class JsonRpcTest { ) } + @Test + fun `should round trip every JSONRPCMessage subtype polymorphically`() { + val messages = listOf( + JSONRPCRequest(id = "request-1", method = "tools/list"), + JSONRPCNotification(method = "notifications/log"), + JSONRPCResponse(id = RequestId("response-1"), result = EmptyResult()), + JSONRPCError( + id = RequestId("error-1"), + error = RPCError(RPCError.ErrorCode.INTERNAL_ERROR, "Internal error"), + ), + JSONRPCEmptyMessage, + ) + + messages.forEach { message -> + val encoded = McpJson.encodeToString(message) + McpJson.decodeFromString(encoded) shouldBe message + } + } + @Test fun `should create JSONRPCRequest with string ID`() { val params = buildJsonObject {