From 9429c40bb07ede969b63ba85b379bedf543475dc Mon Sep 17 00:00:00 2001 From: remote-controlled-man Date: Fri, 18 Sep 2026 14:04:00 +0800 Subject: [PATCH] Fix response result deserialization to use the original request method (#601) An incoming JSON-RPC response's result was decoded by JSON shape alone, so a response whose payload matches another type's shape (e.g. a tasks/result payload shaped like CallToolResult) was completed with the wrong runtime type and failed at the erased generic cast in Protocol.request. - Capture the raw result JSON while decoding JSONRPCResponse on the wire path, keeping the public JSONRPCResponse API unchanged. - At the request/response correlation point, decode the raw result with the deserializer declared by the original request's method, falling back to the shape-decoded result for custom methods, programmatically constructed responses, and task-augmented calls returning CreateTaskResult. - Rewrite the SSE replay message id in the raw JSON before decoding in StreamableHttpClientTransport instead of copying the decoded message, so the raw result survives. - Results whose shape matches no known type now surface as GetTaskPayloadResult (the raw-payload result type) instead of failing the whole message decode. --- .../client/StreamableHttpClientTransport.kt | 36 ++- ...StreamableHttpClientTransportReplayTest.kt | 93 +++++++ .../kotlin/sdk/shared/Protocol.kt | 32 ++- .../kotlin/sdk/types/jsonRpc.kt | 13 + .../kotlin/sdk/types/serializers.kt | 89 ++++++- .../ProtocolResultDeserializationTest.kt | 239 ++++++++++++++++++ 6 files changed, 488 insertions(+), 14 deletions(-) create mode 100644 kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/streamable/http/StreamableHttpClientTransportReplayTest.kt create mode 100644 kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolResultDeserializationTest.kt diff --git a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt index 6df352abe..3e4c62e2a 100644 --- a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt +++ b/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(element) + } + val rewritten = JsonObject(jsonObject + ("id" to McpJson.encodeToJsonElement(replayMessageId))) + return McpJson.decodeFromJsonElement(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(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(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) diff --git a/kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/streamable/http/StreamableHttpClientTransportReplayTest.kt b/kotlin-sdk-client/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/streamable/http/StreamableHttpClientTransportReplayTest.kt new file mode 100644 index 000000000..ea2debf53 --- /dev/null +++ b/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( + 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() + } +} 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..b7bc037ff 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 @@ -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) } 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..427440f8a 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,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 } // ============================================================================ 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..02e7830e3 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 @@ -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> 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? = + 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 { + 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() diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolResultDeserializationTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolResultDeserializationTest.kt new file mode 100644 index 000000000..39035b727 --- /dev/null +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolResultDeserializationTest.kt @@ -0,0 +1,239 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.kotest.matchers.shouldBe +import io.kotest.matchers.types.shouldBeInstanceOf +import io.modelcontextprotocol.kotlin.sdk.types.CallToolRequest +import io.modelcontextprotocol.kotlin.sdk.types.CallToolRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult +import io.modelcontextprotocol.kotlin.sdk.types.CreateTaskResult +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.JSONRPCResponse +import io.modelcontextprotocol.kotlin.sdk.types.ListToolsResult +import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import io.modelcontextprotocol.kotlin.sdk.types.PingRequest +import io.modelcontextprotocol.kotlin.sdk.types.RequestId +import io.modelcontextprotocol.kotlin.sdk.types.TaskMetadata +import io.modelcontextprotocol.kotlin.sdk.types.TaskStatus +import kotlinx.coroutines.async +import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.buildJsonArray +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.decodeFromJsonElement +import kotlinx.serialization.json.encodeToJsonElement +import kotlinx.serialization.json.put +import kotlin.test.BeforeTest +import kotlin.test.Test + +/** + * Regression tests for https://github.com/modelcontextprotocol/kotlin-sdk/issues/601: + * an incoming response must be deserialized according to the original request's method, + * not by guessing the runtime type from the JSON shape of the result. + */ +class ProtocolResultDeserializationTest { + private lateinit var protocol: TestProtocol + private lateinit var transport: RecordingTransport + + @BeforeTest + fun setUp() { + protocol = TestProtocol() + transport = RecordingTransport() + } + + /** + * Decodes a raw wire response exactly like the production transports do + * (`McpJson.decodeFromString(...)`), so the test exercises the real + * polymorphic wire decoding plus the production request/response correlation path. + */ + private fun wireResponse(id: RequestId, result: JsonObject): JSONRPCMessage { + val wire = buildJsonObject { + put("jsonrpc", "2.0") + put("id", McpJson.encodeToJsonElement(id)) + put("result", result) + } + return McpJson.decodeFromJsonElement(wire) + } + + private val callToolPayload: JsonObject = buildJsonObject { + put( + "content", + buildJsonArray { + add( + buildJsonObject { + put("type", "text") + put("text", "task output") + }, + ) + }, + ) + put("isError", false) + } + + private val taskStatePayload: JsonObject = buildJsonObject { + put("taskId", "task-42") + put("status", "completed") + put("createdAt", "2026-01-01T00:00:00Z") + put("lastUpdatedAt", "2026-01-01T00:01:00Z") + } + + @Test + fun `should decode tasks result response as GetTaskPayloadResult even when payload matches another shape`() = + runTest { + protocol.connect(transport) + + val inFlight = async { + protocol.request( + GetTaskPayloadRequest(GetTaskPayloadRequestParams(taskId = "task-42")), + ) + } + val sent = transport.awaitRequest() + + // A tasks/result response for a task-augmented tools/call carries a CallToolResult-shaped + // payload; shape-based decoding resolves it as CallToolResult instead of GetTaskPayloadResult. + transport.deliver(wireResponse(sent.id, callToolPayload)) + + val result = inFlight.await() + result.json shouldBe callToolPayload + } + + @Test + fun `should correlate tasks get and tasks result responses whose payloads share the same shape`() = runTest { + protocol.connect(transport) + + val taskGet = async { + protocol.request(GetTaskRequest(GetTaskRequestParams(taskId = "task-42"))) + } + val taskPayload = async { + protocol.request( + GetTaskPayloadRequest(GetTaskPayloadRequestParams(taskId = "task-42")), + ) + } + val getSent = transport.awaitRequest() + val payloadSent = transport.awaitRequest() + + // Respond in reverse order; both responses carry the exact same GetTaskResult-shaped object, + // which is a valid tasks/result payload (e.g. the payload of a task-augmented tasks/cancel). + transport.deliver(wireResponse(payloadSent.id, taskStatePayload)) + transport.deliver(wireResponse(getSent.id, taskStatePayload)) + + val payload = taskPayload.await() + payload.json shouldBe taskStatePayload + + val get = taskGet.await() + get.taskId shouldBe "task-42" + get.status shouldBe TaskStatus.Completed + } + + @Test + fun `should decode tasks result response whose payload matches no known result shape`() = runTest { + protocol.connect(transport) + + val inFlight = async { + protocol.request( + GetTaskPayloadRequest(GetTaskPayloadRequestParams(taskId = "task-42")), + ) + } + val sent = transport.awaitRequest() + + val unknownPayload = buildJsonObject { put("customField", "custom-value") } + transport.deliver(wireResponse(sent.id, unknownPayload)) + + val result = inFlight.await() + result.json shouldBe unknownPayload + } + + @Test + fun `should decode empty result for ping`() = runTest { + protocol.connect(transport) + + val inFlight = async { protocol.request(PingRequest()) } + val sent = transport.awaitRequest() + + transport.deliver(wireResponse(sent.id, buildJsonObject { })) + + inFlight.await() + } + + @Test + fun `should decode call tool result`() = runTest { + protocol.connect(transport) + + val inFlight = async { + protocol.request(CallToolRequest(CallToolRequestParams(name = "echo"))) + } + val sent = transport.awaitRequest() + + transport.deliver(wireResponse(sent.id, callToolPayload)) + + val result = inFlight.await() + result.isError shouldBe false + } + + @Test + fun `should decode task-augmented call tool response as CreateTaskResult`() = runTest { + protocol.connect(transport) + + val inFlight = async { + protocol.request( + CallToolRequest( + CallToolRequestParams(name = "slow-tool", arguments = null, task = TaskMetadata()), + ), + ) + } + val sent = transport.awaitRequest() + + transport.deliver( + wireResponse( + sent.id, + buildJsonObject { + put( + "task", + buildJsonObject { + put("taskId", "task-7") + put("status", "working") + put("createdAt", "2026-01-01T00:00:00Z") + put("lastUpdatedAt", "2026-01-01T00:00:00Z") + }, + ) + }, + ), + ) + + val result = inFlight.await() + result.task.taskId shouldBe "task-7" + result.task.status shouldBe TaskStatus.Working + } + + @Test + fun `should keep shape-decoded result on the message for direct transport consumers`() { + val message = wireResponse( + RequestId("req-1"), + buildJsonObject { + put( + "tools", + buildJsonArray { + add( + buildJsonObject { + put("name", "echo") + put( + "inputSchema", + buildJsonObject { put("type", "object") }, + ) + }, + ) + }, + ) + }, + ) + + message.shouldBeInstanceOf() + message.result.shouldBeInstanceOf() + } +}