diff --git a/integration-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientTest.kt b/integration-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientTest.kt index 5dfcf61..54d25ee 100644 --- a/integration-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientTest.kt +++ b/integration-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/ClientTest.kt @@ -24,6 +24,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest import io.modelcontextprotocol.kotlin.sdk.types.InitializeResult import io.modelcontextprotocol.kotlin.sdk.types.IntegerSchema import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION @@ -52,6 +53,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.UntitledSingleSelectEnumSchema import io.modelcontextprotocol.kotlin.sdk.types.UrlElicitationRequiredException import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.TimeoutCancellationException +import kotlinx.coroutines.awaitCancellation import kotlinx.coroutines.cancel import kotlinx.coroutines.delay import kotlinx.coroutines.joinAll @@ -71,6 +73,7 @@ import kotlin.test.assertIs import kotlin.test.assertNull import kotlin.test.assertTrue import kotlin.test.fail +import kotlin.time.Duration.Companion.seconds class ClientTest { @Test @@ -543,6 +546,15 @@ class ClientTest { val serverSession = serverSessionResult.await() + // Register the tap after createSession so it chains onto the session's subscription instead + // of pre-empting the handshake. + val cancelledOnWire = CompletableDeferred() + serverTransport.onMessage { message -> + if (message is JSONRPCNotification && message.method == Method.Defined.NotificationsCancelled.value) { + cancelledOnWire.complete(Unit) + } + } + serverSession.setRequestHandler(Method.Defined.ResourcesList) { _, _ -> // Simulate delay def.complete(Unit) @@ -568,6 +580,7 @@ class ClientTest { def.await() runCatching { job.cancel("Cancelled by test") } defCancel.await() + cancelledOnWire.await() defTimeOut.await() } @@ -606,17 +619,9 @@ class ClientTest { val serverSession = serverSessionResult.await() serverSession.setRequestHandler(Method.Defined.ResourcesList) { _, _ -> - // Simulate a delayed response - // Wait ~100ms unless canceled - try { - withTimeout(100L) { - // Just delay here, if timeout is 0 on the client side, this won't return in time - delay(100) - } - } catch (_: Exception) { - // If aborted, just rethrow or return early - } - ListResourcesResult(resources = emptyList()) + // Never respond: the handler runs on a real dispatcher while the 1 ms timeout runs on the + // virtual test clock, so a real-time delay could race it. Parking keeps the test deterministic. + awaitCancellation() } // Request with 1 msec timeout should fail immediately @@ -895,8 +900,8 @@ class ClientTest { ), ) - // Track notifications - var rootListChangedNotificationReceived = false + // The handler runs on Dispatchers.Default after concurrent dispatch is enabled; await the deferred. + val rootListChangedReceived = CompletableDeferred() val serverSessionResult = CompletableDeferred() @@ -915,16 +920,13 @@ class ClientTest { serverSession.setNotificationHandler( Method.Defined.NotificationsRootsListChanged, ) { - rootListChangedNotificationReceived = true + rootListChangedReceived.complete(Unit) CompletableDeferred(Unit) } client.sendRootsListChanged() - assertTrue( - rootListChangedNotificationReceived, - "Notification should be sent when sendRootsListChanged is called", - ) + withTimeout(5.seconds) { rootListChangedReceived.await() } } @Test diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ProtocolConcurrencyIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ProtocolConcurrencyIntegrationTest.kt new file mode 100644 index 0000000..d8d64ec --- /dev/null +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ProtocolConcurrencyIntegrationTest.kt @@ -0,0 +1,306 @@ +package io.modelcontextprotocol.kotlin.sdk.integration + +import io.kotest.matchers.shouldBe +import io.modelcontextprotocol.kotlin.sdk.ExperimentalMcpApi +import io.modelcontextprotocol.kotlin.sdk.client.Client +import io.modelcontextprotocol.kotlin.sdk.client.ClientOptions +import io.modelcontextprotocol.kotlin.sdk.server.Server +import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions +import io.modelcontextprotocol.kotlin.sdk.server.ServerSession +import io.modelcontextprotocol.kotlin.sdk.shared.InMemoryTransport +import io.modelcontextprotocol.kotlin.sdk.shared.currentRequestHandlerExtra +import io.modelcontextprotocol.kotlin.sdk.testing.ChannelTransport +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.ClientCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageRequest +import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageResult +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest +import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse +import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION +import io.modelcontextprotocol.kotlin.sdk.types.ListRootsRequest +import io.modelcontextprotocol.kotlin.sdk.types.ListRootsResult +import io.modelcontextprotocol.kotlin.sdk.types.ListToolsRequest +import io.modelcontextprotocol.kotlin.sdk.types.ListToolsResult +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.RequestId +import io.modelcontextprotocol.kotlin.sdk.types.Role +import io.modelcontextprotocol.kotlin.sdk.types.SamplingMessage +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import io.modelcontextprotocol.kotlin.sdk.types.toJSON +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeout +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.CopyOnWriteArrayList +import java.util.concurrent.atomic.AtomicReference +import kotlin.coroutines.cancellation.CancellationException +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.time.Duration.Companion.seconds + +/** + * End-to-end tests for concurrent inbound dispatch enabled after the MCP handshake. + * + * Handlers run on a real [kotlinx.coroutines.Dispatchers.Default], so these use [runBlocking] with + * explicit [withTimeout] guards. A virtual test clock would false-fire timeouts against cross-thread + * real dispatchers. + */ +@OptIn(ExperimentalMcpApi::class) +class ProtocolConcurrencyIntegrationTest { + + private fun rootsClient(): Client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions( + capabilities = ClientCapabilities(roots = ClientCapabilities.Roots(listChanged = null)), + ), + ) + + private fun toolsServer(): Server = Server( + serverInfo = Implementation(name = "test server", version = "1.0"), + options = ServerOptions( + capabilities = ServerCapabilities(tools = ServerCapabilities.Tools(listChanged = null)), + ), + ) + + // ChannelTransport delivers inbound messages as soon as its read loop starts, which is before + // start() marks it operational. Creating the session first keeps the client's initialize from + // landing on a transport that would reject the session's response. + private suspend fun connect( + client: Client, + clientTransport: ChannelTransport, + server: Server, + serverTransport: ChannelTransport, + ): ServerSession { + val serverSession = server.createSession(serverTransport) + client.connect(clientTransport) + return serverSession + } + + // The motivating deadlock over real serial read loops: a server tool handler makes a + // server-to-client request (roots/list) while still handling the client's tools/call. + @Test + fun `server tool handler can call roots list mid-handling without deadlock`() = runBlocking { + val client = rootsClient() + val server = toolsServer() + val (clientTransport, serverTransport) = ChannelTransport.createLinkedPair() + + val serverSession = connect(client, clientTransport, server, serverTransport) + + // Register the tools/call handler directly on the session so it can call back to the client + // mid-handling. Before concurrent dispatch this would deadlock: the server read loop is + // blocked by the tool handler and can never deliver the roots/list response. + serverSession.setRequestHandler(Method.Defined.ToolsCall) { _, _ -> + val roots = serverSession.request(ListRootsRequest()) + CallToolResult(content = listOf(TextContent(text = "roots=${roots.roots.size}"))) + } + + val result = withTimeout(10.seconds) { + client.callTool(CallToolRequest(CallToolRequestParams("needs-roots"))) + } + val text = (result as CallToolResult).content.filterIsInstance().single().text + text shouldBe "roots=0" + + client.close() + server.close() + } + + // Client-side mirror: a client sampling handler makes a client-to-server request (tools/list) + // while still handling the server's sampling/createMessage. + @Test + fun `client sampling handler can call the server mid-handling without deadlock`() = runBlocking { + val client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions( + capabilities = ClientCapabilities(sampling = ClientCapabilities.Sampling()), + ), + ) + val server = toolsServer() + val (clientTransport, serverTransport) = ChannelTransport.createLinkedPair() + + client.setRequestHandler(Method.Defined.SamplingCreateMessage) { _, _ -> + val tools = client.request(ListToolsRequest()) + CreateMessageResult( + role = Role.Assistant, + content = TextContent(text = "tools=${tools.tools.size}"), + model = "test-model", + ) + } + + val serverSession = connect(client, clientTransport, server, serverTransport) + + val result = withTimeout(10.seconds) { + serverSession.request( + CreateMessageRequest( + CreateMessageRequestParams( + maxTokens = 100, + messages = listOf(SamplingMessage(role = Role.User, content = TextContent(text = "hi"))), + ), + ), + ) + } + val text = result.content.filterIsInstance().single().text + text shouldBe "tools=0" + + client.close() + server.close() + } + + // Nested-cancel cascade: cancelling the client's tools/call must reach the server tool handler + // over the wire and cascade into its nested roots/list request, cancelling the parked + // client-side roots handler. No tools/call response is ever sent. + @Test + fun `cancelling a tool call cascades into its nested roots list request`() = runBlocking { + val client = rootsClient() + val server = toolsServer() + val (clientTransport, serverTransport) = ChannelTransport.createLinkedPair() + + val capturedToolCallId = AtomicReference(null) + val toolHandlerCancelled = CompletableDeferred() + val rootsHandlerEntered = CompletableDeferred() + val rootsHandlerCancelled = CompletableDeferred() + val parkForever = CompletableDeferred() + val unexpectedToolCallResponse = CompletableDeferred() + + // Tap the server->client direction (client transport inbound) to prove no tools/call + // response is ever delivered once the call is cancelled. + clientTransport.onMessage { message -> + if (message is JSONRPCResponse && message.id == capturedToolCallId.get()) { + unexpectedToolCallResponse.complete(Unit) + } + } + + val serverSession = connect(client, clientTransport, server, serverTransport) + + // Client roots handler parks until cancelled and records the cancellation. + client.setRequestHandler(Method.Defined.RootsList) { _, _ -> + rootsHandlerEntered.complete(Unit) + try { + parkForever.await() + } catch (e: CancellationException) { + rootsHandlerCancelled.complete(e) + throw e + } + } + + // Tool handler makes a nested server->client request and records its own cancellation. + serverSession.setRequestHandler(Method.Defined.ToolsCall) { _, _ -> + val extra = currentRequestHandlerExtra()!! + capturedToolCallId.set(extra.requestId) + try { + extra.sendRequest(ListRootsRequest()) + } catch (e: CancellationException) { + toolHandlerCancelled.complete(e) + throw e + } + CallToolResult(content = listOf(TextContent(text = "unreachable"))) + } + + val call = launch { client.callTool(CallToolRequest(CallToolRequestParams("nested"))) } + + withTimeout(10.seconds) { rootsHandlerEntered.await() } + call.cancel(CancellationException("user aborted")) + + withTimeout(10.seconds) { toolHandlerCancelled.await() } + withTimeout(10.seconds) { rootsHandlerCancelled.await() } + + // The server suppresses the response for a cancelled request, so none is ever sent. + assertFalse( + unexpectedToolCallResponse.isCompleted, + "no tools/call response should be delivered after the call was cancelled", + ) + + client.close() + server.close() + } + + // Init gate plus context element in a high-level handler: initialize and tools/list are + // pipelined before notifications/initialized and answered serially in order. After the + // handshake a high-level addTool handler can observe currentRequestHandlerExtra(). + @Test + fun `pipelined initialize and tools list are processed serially then context element is available`(): Unit = + runBlocking { + val server = toolsServer() + server.addTool(name = "ctx-tool", description = "returns its request id") { + val extra = currentRequestHandlerExtra() + checkNotNull(extra) { "currentRequestHandlerExtra() must be non-null inside a tool handler" } + CallToolResult(content = listOf(TextContent(text = extra.requestId.asText()))) + } + + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + + val arrivalOrder = CopyOnWriteArrayList() + val pending = ConcurrentHashMap>() + clientTransport.onMessage { message -> + val id = when (message) { + is JSONRPCResponse -> message.id + is JSONRPCError -> message.id + else -> null + } + if (id != null) { + arrivalOrder.add(id) + pending[id]?.complete(message) + } + } + + // Drive the client side manually over the raw transport. + server.createSession(serverTransport) + + val initRequest = InitializeRequest( + InitializeRequestParams( + protocolVersion = LATEST_PROTOCOL_VERSION, + capabilities = ClientCapabilities(), + clientInfo = Implementation(name = "raw-client", version = "1.0"), + ), + ).toJSON() + val toolsListRequest = ListToolsRequest().toJSON() + + val initAnswered = pending.expect(initRequest.id) + val toolsListAnswered = pending.expect(toolsListRequest.id) + + // Pipelined before notifications/initialized: both handled in the serial dispatch phase. + clientTransport.send(initRequest) + clientTransport.send(toolsListRequest) + + withTimeout(10.seconds) { initAnswered.await() } + withTimeout(10.seconds) { toolsListAnswered.await() } + + // Deterministic serial ordering: initialize answered before tools/list. + assertEquals( + listOf(initRequest.id, toolsListRequest.id), + arrivalOrder.toList(), + "pipelined requests must be answered in arrival order", + ) + + // Complete the handshake, then call the high-level tool. + clientTransport.send(InitializedNotification().toJSON()) + + val toolCallRequest = CallToolRequest(CallToolRequestParams("ctx-tool")).toJSON() + val toolCallAnswered = pending.expect(toolCallRequest.id) + clientTransport.send(toolCallRequest) + + val toolCallResponse = withTimeout(10.seconds) { toolCallAnswered.await() } as JSONRPCResponse + val text = (toolCallResponse.result as CallToolResult).content.filterIsInstance().single().text + // The handler echoed its request id, proving currentRequestHandlerExtra() was available. + text shouldBe toolCallRequest.id.asText() + } + + private fun ConcurrentHashMap>.expect( + id: RequestId, + ): CompletableDeferred = CompletableDeferred().also { put(id, it) } + + private fun RequestId.asText(): String = when (this) { + is RequestId.StringId -> value + is RequestId.NumberId -> value.toString() + } +} diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ClientConnectionLoggingTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ClientConnectionLoggingTest.kt index 81a3b41..a1cf797 100644 --- a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ClientConnectionLoggingTest.kt +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ClientConnectionLoggingTest.kt @@ -1,5 +1,7 @@ package io.modelcontextprotocol.kotlin.sdk.server +import io.kotest.assertions.nondeterministic.continually +import io.kotest.assertions.nondeterministic.eventually import io.kotest.matchers.collections.shouldBeEmpty import io.kotest.matchers.collections.shouldHaveSize import io.kotest.matchers.shouldBe @@ -17,6 +19,9 @@ import kotlinx.serialization.json.JsonPrimitive import org.junit.jupiter.api.Test import org.junit.jupiter.params.ParameterizedTest import org.junit.jupiter.params.provider.EnumSource +import java.util.concurrent.CopyOnWriteArrayList +import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.seconds class ClientConnectionLoggingTest : AbstractServerFeaturesTest() { @@ -56,7 +61,7 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() { @ParameterizedTest @EnumSource(LoggingLevel::class) fun `notification should filter logging messages below level`(minLevel: LoggingLevel): Unit = runBlocking { - val receivedMessages = mutableListOf() + val receivedMessages = CopyOnWriteArrayList() client.setNotificationHandler(Method.Defined.NotificationsMessage) { receivedMessages.add(it) CompletableDeferred(Unit) @@ -79,8 +84,10 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() { client.callTool(CallToolRequest(CallToolRequestParams("test-notification-filtered"))) + // Handlers dispatch on Dispatchers.Default after the handshake, so await the expected count + // before asserting order. val expectedLevels = LoggingLevel.entries.filter { it >= minLevel } - receivedMessages shouldHaveSize expectedLevels.size + eventually(5.seconds) { receivedMessages shouldHaveSize expectedLevels.size } receivedMessages.map { it.params.level } shouldBe expectedLevels } @@ -115,7 +122,7 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() { @ParameterizedTest @EnumSource(LoggingLevel::class) fun `sendLoggingMessage should filter messages below level`(minLevel: LoggingLevel): Unit = runBlocking { - val receivedMessages = mutableListOf() + val receivedMessages = CopyOnWriteArrayList() client.setNotificationHandler(Method.Defined.NotificationsMessage) { receivedMessages.add(it) CompletableDeferred(Unit) @@ -138,14 +145,16 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() { client.callTool(CallToolRequest(CallToolRequestParams("test-logging-level"))) + // Handlers dispatch on Dispatchers.Default after the handshake, so await the expected count + // before asserting order. val expectedLevels = LoggingLevel.entries.filter { it >= minLevel } - receivedMessages shouldHaveSize expectedLevels.size + eventually(5.seconds) { receivedMessages shouldHaveSize expectedLevels.size } receivedMessages.map { it.params.level } shouldBe expectedLevels } @Test fun `sendLoggingMessage should send no messages when level is set to highest`(): Unit = runBlocking { - val receivedMessages = mutableListOf() + val receivedMessages = CopyOnWriteArrayList() client.setNotificationHandler(Method.Defined.NotificationsMessage) { receivedMessages.add(it) CompletableDeferred(Unit) @@ -167,6 +176,7 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() { } client.callTool(CallToolRequest(CallToolRequestParams("test-logging-highest"))) - receivedMessages.shouldBeEmpty() + // Handlers dispatch asynchronously, so hold for a window to confirm nothing ever arrives. + continually(500.milliseconds) { receivedMessages.shouldBeEmpty() } } } diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt index 5031a9c..82cf91a 100644 --- a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt @@ -65,6 +65,8 @@ class ServerSessionInitializeTest { assertEquals("test-client", session.clientVersion?.name) } + // Both initialize requests arrive before notifications/initialized, in the serial dispatch + // phase, so processing and response order stay deterministic. @Test fun `should reject duplicate initialize request`() = runTest { val session = createSession() diff --git a/kotlin-sdk-client/api/kotlin-sdk-client.api b/kotlin-sdk-client/api/kotlin-sdk-client.api index 1493a30..d856f5b 100644 --- a/kotlin-sdk-client/api/kotlin-sdk-client.api +++ b/kotlin-sdk-client/api/kotlin-sdk-client.api @@ -52,8 +52,8 @@ public final class io/modelcontextprotocol/kotlin/sdk/client/ClientKt { public final class io/modelcontextprotocol/kotlin/sdk/client/ClientOptions : io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions { public fun ()V - public fun (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;Z)V - public synthetic fun (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;ZLkotlin/coroutines/CoroutineContext;)V + public synthetic fun (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;ZLkotlin/coroutines/CoroutineContext;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun getCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities; } diff --git a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt index 59f2b7e..3352104 100644 --- a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt +++ b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/Client.kt @@ -73,11 +73,13 @@ import kotlinx.collections.immutable.minus import kotlinx.collections.immutable.persistentMapOf import kotlinx.collections.immutable.toPersistentSet import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers import kotlinx.serialization.SerializationException import kotlinx.serialization.json.JsonElement import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive import kotlinx.serialization.json.buildJsonArray +import kotlin.coroutines.CoroutineContext import kotlin.coroutines.cancellation.CancellationException private val logger = KotlinLogging.logger {} @@ -87,11 +89,16 @@ private val logger = KotlinLogging.logger {} * * @property capabilities The capabilities this client supports. * @param enforceStrictCapabilities Whether to strictly enforce capabilities when interacting with the server. + * @param handlerCoroutineContext Coroutine context for inbound handlers. See [ProtocolOptions.handlerCoroutineContext]. */ public class ClientOptions( public val capabilities: ClientCapabilities = ClientCapabilities(), enforceStrictCapabilities: Boolean = true, -) : ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities) + handlerCoroutineContext: CoroutineContext = Dispatchers.Default, +) : ProtocolOptions( + enforceStrictCapabilities = enforceStrictCapabilities, + handlerCoroutineContext = handlerCoroutineContext, +) /** * Initializes and connects an MCP client using the provided clientInfo [Implementation], client options, @@ -216,6 +223,7 @@ public open class Client(private val clientInfo: Implementation, options: Client serverInstructions = result.instructions notification(InitializedNotification()) + enableConcurrentDispatch() } catch (error: Throwable) { logger.error(error) { "Failed to initialize client: ${error.message}" } close() diff --git a/kotlin-sdk-core/api/kotlin-sdk-core.api b/kotlin-sdk-core/api/kotlin-sdk-core.api index 1935e9a..28d6737 100644 --- a/kotlin-sdk-core/api/kotlin-sdk-core.api +++ b/kotlin-sdk-core/api/kotlin-sdk-core.api @@ -59,6 +59,7 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol { public abstract fun assertRequestHandlerCapability (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V public final fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public fun connect (Lio/modelcontextprotocol/kotlin/sdk/shared/Transport;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + protected final fun enableConcurrentDispatch ()V public final fun getFallbackNotificationHandler ()Lkotlin/jvm/functions/Function2; public final fun getFallbackRequestHandler ()Lkotlin/jvm/functions/Function3; public final fun getNotificationHandlers ()Ljava/util/Map; @@ -71,6 +72,7 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol { public static synthetic fun notification$default (Lio/modelcontextprotocol/kotlin/sdk/shared/Protocol;Lio/modelcontextprotocol/kotlin/sdk/types/Notification;Lio/modelcontextprotocol/kotlin/sdk/types/RequestId;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; public fun onClose ()V public fun onError (Ljava/lang/Throwable;)V + protected fun onInitializedNotification ()V public final fun removeNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V public final fun removeRequestHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V public final fun request (Lio/modelcontextprotocol/kotlin/sdk/types/Request;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; @@ -84,15 +86,18 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol { public final class io/modelcontextprotocol/kotlin/sdk/shared/ProtocolKt { public static final field IMPLEMENTATION_NAME Ljava/lang/String; + public static final fun currentRequestHandlerExtra (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public static final fun getDEFAULT_REQUEST_TIMEOUT ()J } public class io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions { - public synthetic fun (ZJILkotlin/jvm/internal/DefaultConstructorMarker;)V - public synthetic fun (ZJLkotlin/jvm/internal/DefaultConstructorMarker;)V + public synthetic fun (ZJLkotlin/coroutines/CoroutineContext;IIILkotlin/jvm/internal/DefaultConstructorMarker;)V + public synthetic fun (ZJLkotlin/coroutines/CoroutineContext;IILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun getEnforceStrictCapabilities ()Z + public final fun getHandlerCoroutineContext ()Lkotlin/coroutines/CoroutineContext; public final fun getTimeout-UwyO8pc ()J public final fun setEnforceStrictCapabilities (Z)V + public final fun setHandlerCoroutineContext (Lkotlin/coroutines/CoroutineContext;)V public final fun setTimeout-LRDsOJo (J)V } @@ -107,8 +112,16 @@ public final class io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferKt { public static final fun serializeMessage (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;)Ljava/lang/String; } -public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra { - public fun ()V +public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra : kotlin/coroutines/AbstractCoroutineContextElement { + public static final field Key Lio/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra$Key; + public final fun getMethod ()Lio/modelcontextprotocol/kotlin/sdk/types/Method; + public final fun getRequestId ()Lio/modelcontextprotocol/kotlin/sdk/types/RequestId; + public final fun sendNotification (Lio/modelcontextprotocol/kotlin/sdk/types/Notification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public final fun sendRequest (Lio/modelcontextprotocol/kotlin/sdk/types/Request;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public static synthetic fun sendRequest$default (Lio/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra;Lio/modelcontextprotocol/kotlin/sdk/types/Request;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; +} + +public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra$Key : kotlin/coroutines/CoroutineContext$Key { } public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestOptions : io/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions { 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 50dc356..3515b4a 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 @@ -27,15 +27,33 @@ import kotlinx.atomicfu.AtomicRef import kotlinx.atomicfu.atomic import kotlinx.atomicfu.getAndUpdate import kotlinx.atomicfu.update +import kotlinx.collections.immutable.PersistentList import kotlinx.collections.immutable.PersistentMap +import kotlinx.collections.immutable.persistentListOf import kotlinx.collections.immutable.persistentMapOf import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineName +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Deferred -import kotlinx.coroutines.TimeoutCancellationException -import kotlinx.coroutines.withTimeout +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job +import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.cancel +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.job +import kotlinx.coroutines.launch +import kotlinx.coroutines.sync.Semaphore +import kotlinx.coroutines.sync.withPermit +import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeoutOrNull import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive import kotlinx.serialization.json.encodeToJsonElement +import kotlin.coroutines.AbstractCoroutineContextElement +import kotlin.coroutines.ContinuationInterceptor +import kotlin.coroutines.CoroutineContext import kotlin.coroutines.cancellation.CancellationException import kotlin.time.Duration import kotlin.time.Duration.Companion.seconds @@ -53,17 +71,26 @@ public typealias ProgressCallback = (Progress) -> Unit /** * Additional initialization options. * - * @property enforceStrictCapabilities whether to restrict emitted requests to only those that the - * remote side has indicated it can handle through its advertised capabilities. - * This does NOT affect checking of _local_ side capabilities, as it is considered a logic error - * to mis-specify those. - * Currently defaults to `false` for backwards compatibility with SDK versions that did not advertise - * capabilities correctly; in the future, this will default to `true`. + * [handlerCoroutineContext] is read once when [Protocol.connect] attaches a transport and stays + * fixed until the next [Protocol.connect]. + * + * @property enforceStrictCapabilities whether to restrict emitted requests to only those the remote + * side advertised support for. Local-side capability checks are unaffected. Defaults to `false` for + * backward compatibility with SDK versions that did not advertise capabilities correctly. * @property timeout default timeout for outgoing requests + * @property handlerCoroutineContext coroutine context used to run inbound request and notification + * handlers once initialization has completed. Must contain a real dispatching + * [kotlin.coroutines.ContinuationInterceptor]. `Dispatchers.Unconfined` and unconfined test + * dispatchers are unsupported, since handler resumptions would then run on the transport read loop + * and reintroduce head-of-line blocking. Any [kotlinx.coroutines.Job] in this context is ignored. */ public open class ProtocolOptions( public var enforceStrictCapabilities: Boolean = false, public var timeout: Duration = DEFAULT_REQUEST_TIMEOUT, + public var handlerCoroutineContext: CoroutineContext = Dispatchers.Default, + // Internal handler-concurrency bounds, read once at connect and not part of the public surface. + internal var maxConcurrentHandlers: Int = DEFAULT_MAX_CONCURRENT_HANDLERS, + internal var maxInFlightHandlers: Int = DEFAULT_MAX_IN_FLIGHT_HANDLERS, ) /** @@ -71,6 +98,12 @@ public open class ProtocolOptions( */ public val DEFAULT_REQUEST_TIMEOUT: Duration = 60.seconds +/** Default cap on concurrently executing inbound handlers per connection. */ +internal const val DEFAULT_MAX_CONCURRENT_HANDLERS: Int = 64 + +/** Default cap on launched-but-not-completed inbound handlers per connection. */ +internal const val DEFAULT_MAX_IN_FLIGHT_HANDLERS: Int = 256 + /** * Options that can be given per request. * @@ -121,7 +154,7 @@ public class RequestOptions( override fun hashCode(): Int { var result = super.hashCode() - result = 31 * result + (onProgress?.hashCode() ?: 0) + result = 31 * result + onProgress.hashCode() result = 31 * result + timeout.hashCode() return result } @@ -132,11 +165,70 @@ public class RequestOptions( /** * Extra data given to request handlers. + * + * One instance exists per in-flight inbound request. It is passed as the second parameter to + * low-level handlers registered via [Protocol.setRequestHandler] **and** installed into the + * handler coroutine's [CoroutineContext], so code nested anywhere under the handler can reach it + * via [currentRequestHandlerExtra] without parameter plumbing. + * + * @property requestId the JSON-RPC id of the request being handled + * @property method the request's method */ -public class RequestHandlerExtra +public class RequestHandlerExtra internal constructor( + public val requestId: RequestId, + public val method: Method, + private val protocol: Protocol, + internal val capturedTransport: Transport, +) : AbstractCoroutineContextElement(Key) { + + /** Context key for retrieving the extra from a handler coroutine's context. */ + public companion object Key : CoroutineContext.Key + + /** + * Sends a notification tagged with `relatedRequestId = `[requestId], letting transports + * associate it with the request being handled (e.g. route it onto the request's SSE stream). + */ + public suspend fun sendNotification(notification: Notification) { + protocol.notification(notification, relatedRequestId = requestId) + } + + /** + * Sends a request to the peer tagged with `relatedRequestId = `[requestId]. + * + * Mirrors [Protocol.request]; cancellation of the handler coroutine propagates into this call. + */ + public suspend fun sendRequest(request: Request, options: RequestOptions? = null): T { + val base = options ?: RequestOptions() + return protocol.request( + request, + RequestOptions( + relatedRequestId = requestId, + resumptionToken = base.resumptionToken, + onResumptionToken = base.onResumptionToken, + onProgress = base.onProgress, + timeout = base.timeout, + ), + ) + } +} + +/** The [RequestHandlerExtra] of the MCP request being handled by the current coroutine, or `null` outside a handler. */ +public suspend fun currentRequestHandlerExtra(): RequestHandlerExtra? = currentCoroutineContext()[RequestHandlerExtra] internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) } +/** Cap on how many recently-cancelled outbound request ids are remembered to quiet late responses/progress. */ +private const val CANCELLED_REQUEST_IDS_REMEMBERED = 256 + +// Bypass set: exempt from both bound tiers so the connection cannot self-block. +// Responses are not listed — they are handled inline before dispatch. +private val CONTROL_METHODS: Set = setOf( + Method.Defined.Ping.value, + Method.Defined.NotificationsCancelled.value, + Method.Defined.NotificationsProgress.value, + Method.Defined.NotificationsInitialized.value, +) + /** * Implements MCP protocol framing on top of a pluggable transport, including * features like request/response linking, notifications, and progress. @@ -144,9 +236,11 @@ internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) } * @property options protocol-level configuration; `null` falls back to defaults */ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptions?) { + private val connectionRef: AtomicRef = atomic(null) + /** The active transport, or `null` if not connected. */ - public var transport: Transport? = null - private set + public val transport: Transport? + get() = connectionRef.value?.transport private val _requestHandlers: AtomicRef RequestResult?>> = @@ -184,6 +278,21 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio public val progressHandlers: Map get() = _progressHandlers.value + private val recentlyCancelledRequestIds: AtomicRef> = + atomic(persistentListOf()) + + /** Remembers a locally-cancelled outbound request id so its late response/progress is quieted. */ + private fun rememberCancelledRequestId(id: RequestId) { + recentlyCancelledRequestIds.update { current -> + val appended = current.adding(id) + if (appended.size > CANCELLED_REQUEST_IDS_REMEMBERED) appended.removingAt(0) else appended + } + } + + /** Whether [id] was recently cancelled locally (see [rememberCancelledRequestId]). */ + private fun isRecentlyCancelled(id: RequestId?): Boolean = + id != null && recentlyCancelledRequestIds.value.contains(id) + /** * Callback for when the connection is closed for any reason. * @@ -213,6 +322,11 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio public var fallbackNotificationHandler: (suspend (notification: JSONRPCNotification) -> Unit)? = null init { + setNotificationHandler(Method.Defined.NotificationsCancelled) { notification -> + handleCancelledNotification(notification) + COMPLETED + } + setNotificationHandler(Method.Defined.NotificationsProgress) { notification -> onProgress(notification) COMPLETED @@ -223,15 +337,64 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } } + /** + * Enables concurrent dispatch of inbound messages on the current connection. + * + * Called by subclasses when the MCP initialization phase completes (the client after the + * initialize exchange; the server session on receiving `notifications/initialized`). Until + * then, inbound messages are handled inline in arrival order. Reset by [connect]. + */ + protected fun enableConcurrentDispatch() { + val connection = connectionRef.value ?: return + connection.concurrentDispatchEnabled.value = true + } + + /** + * Router hook invoked whenever `notifications/initialized` is received, before handler lookup. + * + * Runs regardless of which notification handler is registered for the method, so replacing the + * handler via [setNotificationHandler] cannot disable subclass initialization logic. + */ + protected open fun onInitializedNotification() {} + /** * Attaches to the given transport, starts it, and starts listening for messages. * - * The Protocol object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward. + * The Protocol object assumes ownership of the Transport, replacing any callbacks that have + * already been set, and expects that it is the only user of the Transport instance going forward. + * + * @throws IllegalStateException if this Protocol is already connected; call [close] first. */ public open suspend fun connect(transport: Transport) { - this.transport = transport + val maxConcurrent = options?.maxConcurrentHandlers ?: DEFAULT_MAX_CONCURRENT_HANDLERS + val maxInFlight = options?.maxInFlightHandlers ?: DEFAULT_MAX_IN_FLIGHT_HANDLERS + require(maxConcurrent > 0) { "maxConcurrentHandlers must be positive, but was $maxConcurrent" } + require(maxInFlight >= maxConcurrent) { + "maxInFlightHandlers ($maxInFlight) must be >= maxConcurrentHandlers ($maxConcurrent)" + } + val handlerContext = (options?.handlerCoroutineContext ?: Dispatchers.Default).minusKey(Job) + val interceptor = handlerContext[ContinuationInterceptor] + if (interceptor === Dispatchers.Unconfined || + interceptor?.let { it::class.simpleName?.contains("Unconfined") } == true + ) { + logger.warn { + "handlerCoroutineContext uses an unconfined dispatcher ($interceptor); " + + "handler resumptions will run on the transport read loop and may block unrelated messages" + } + } + val connection = Connection( + transport = transport, + handlerScope = CoroutineScope(SupervisorJob() + handlerContext + CoroutineName("McpProtocol")), + executionSemaphore = Semaphore(maxConcurrent), + maxInFlightHandlers = maxInFlight, + ) + if (!connectionRef.compareAndSet(null, connection)) { + connection.handlerScope.cancel() + error("Protocol is already connected; close() the current connection before connecting a new transport") + } + transport.onClose { - doClose() + doClose(connection) } transport.onError { @@ -239,24 +402,42 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } transport.onMessage { message -> + if (connectionRef.value !== connection) { + logger.trace { "Dropping message received after close: ${message::class.simpleName}" } + return@onMessage + } when (message) { is JSONRPCResponse -> onResponse(message, null) - is JSONRPCRequest -> onRequest(message) - is JSONRPCNotification -> onNotification(message) + is JSONRPCRequest -> dispatchRequest(connection, message) + is JSONRPCNotification -> dispatchNotification(connection, message) is JSONRPCError -> onResponse(null, message) is JSONRPCEmptyMessage -> Unit } } logger.info { "Starting transport" } - transport.start() + try { + transport.start() + } catch (cause: Throwable) { + // Roll back so the Protocol can be reconnected after a failed start; idempotent with + // a transport-initiated doClose (both CAS on the same connection). + if (connectionRef.compareAndSet(connection, null)) { + connection.handlerScope.cancel() + } + throw cause + } } - private fun doClose() { - val handlersToNotify = _responseHandlers.value.values.toList() - _responseHandlers.getAndSet(persistentMapOf()) + private fun doClose(connection: Connection) { + // A stale onClose from a previous transport must not tear down the successor connection. + if (!connectionRef.compareAndSet(connection, null)) { + logger.trace { "Ignoring close signal from a stale transport" } + return + } + val handlersToNotify = _responseHandlers.getAndSet(persistentMapOf()).values.toList() _progressHandlers.getAndSet(persistentMapOf()) - transport = null + connection.inFlightRequestJobs.getAndSet(persistentMapOf()) + connection.handlerScope.cancel() onClose() val error = McpException(RPCError.ErrorCode.CONNECTION_CLOSED, "Connection closed") @@ -265,6 +446,110 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } } + private suspend fun dispatchRequest(connection: Connection, request: JSONRPCRequest) { + if (!connection.concurrentDispatchEnabled.value) { + // Serial phase (before the peer completes MCP initialization): inline, in arrival + // order — the delivering coroutine awaits the handler. + onRequest(request, connection, trackCancellation = false) + return + } + val bypass = request.method in CONTROL_METHODS + if (!bypass && !tryAdmit(connection)) { + rejectOverloadedRequest(connection, request) + return + } + // UNDISPATCHED starts the handler synchronously in the delivering coroutine, preserving + // arrival order, then moves to handlerCoroutineContext at its first suspension. + connection.handlerScope.launch(start = CoroutineStart.UNDISPATCHED) { + // Register before any suspension point so a following notifications/cancelled can find + // and cancel even a handler still parked on admission. + registerInFlight(connection, request.id, coroutineContext.job) + try { + if (bypass) { + onRequest(request, connection, trackCancellation = true) + } else { + connection.executionSemaphore.withPermit { + onRequest(request, connection, trackCancellation = true) + } + } + } finally { + if (!bypass) connection.inFlightCount.decrementAndGet() + } + } + } + + private suspend fun dispatchNotification(connection: Connection, notification: JSONRPCNotification) { + if (notification.method == Method.Defined.NotificationsInitialized.value) { + onInitializedNotification() + } + if (!connection.concurrentDispatchEnabled.value) { + onNotification(notification) + return + } + val bypass = notification.method in CONTROL_METHODS + if (!bypass && !tryAdmit(connection)) { + val dropped = IllegalStateException( + "Dropped notification ${notification.method}: too many in-flight messages", + ) + logger.warn { dropped.message } + onError(dropped) + return + } + connection.handlerScope.launch(start = CoroutineStart.UNDISPATCHED) { + try { + if (bypass) { + onNotification(notification) + } else { + connection.executionSemaphore.withPermit { onNotification(notification) } + } + } finally { + if (!bypass) connection.inFlightCount.decrementAndGet() + } + } + } + + private fun tryAdmit(connection: Connection): Boolean { + while (true) { + val current = connection.inFlightCount.value + if (current >= connection.maxInFlightHandlers) return false + if (connection.inFlightCount.compareAndSet(current, current + 1)) return true + } + } + + private suspend fun rejectOverloadedRequest(connection: Connection, request: JSONRPCRequest) { + logger.warn { + "Rejecting request ${request.method} (id: ${request.id}): " + + "in-flight handler limit (${connection.maxInFlightHandlers}) reached" + } + try { + connection.transport.send( + JSONRPCError( + id = request.id, + error = RPCError( + code = RPCError.ErrorCode.INTERNAL_ERROR, + message = "Server is busy: too many in-flight messages", + ), + ), + ) + } catch (e: CancellationException) { + throw e + } catch (cause: Throwable) { + logger.error(cause) { "Failed to send busy rejection for request id: ${request.id}" } + onError(cause) + } + } + + private fun handleCancelledNotification(notification: CancelledNotification) { + val connection = connectionRef.value ?: return + val requestId = notification.params.requestId + val job = connection.inFlightRequestJobs.value[requestId] + if (job == null) { + logger.trace { "Ignoring cancellation for unknown or already-completed request: $requestId" } + return + } + job.cancel("Cancelled by peer: ${notification.params.reason ?: "unknown"}") + } + private suspend fun onNotification(notification: JSONRPCNotification) { logger.trace { "Received notification: ${notification.method}" } @@ -284,15 +569,16 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } } - private suspend fun onRequest(request: JSONRPCRequest) { + private suspend fun onRequest(request: JSONRPCRequest, connection: Connection, trackCancellation: Boolean) { logger.trace { "Received request: ${request.method} (id: ${request.id})" } val handler = requestHandlers[request.method] ?: fallbackRequestHandler + val capturedTransport = connection.transport if (handler === null) { logger.trace { "No handler found for request: ${request.method}" } try { - transport?.send( + capturedTransport.send( JSONRPCError( id = request.id, error = RPCError( @@ -310,36 +596,77 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio return } - try { - val result = handler(request, RequestHandlerExtra()) - logger.trace { "Request handled successfully: ${request.method} (id: ${request.id})" } + val extra = RequestHandlerExtra( + requestId = request.id, + method = Method.from(request.method), + protocol = this, + capturedTransport = capturedTransport, + ) - transport?.send( - JSONRPCResponse( - id = request.id, - result = result ?: EmptyResult(), - ), - ) + // Registration in the in-flight registry already happened in dispatchRequest's launch + // block (synchronously, pre-suspension). Here we only need the job for suppression checks. + val handlerJob = if (trackCancellation) currentCoroutineContext().job else null + + withContext(extra) { + try { + val result = handler(request, extra) + logger.trace { "Request handled successfully: ${request.method} (id: ${request.id})" } + if (handlerJob?.isCancelled == true) { + logger.trace { "Suppressing response for request cancelled by peer (id: ${request.id})" } + return@withContext + } + capturedTransport.send( + JSONRPCResponse( + id = request.id, + result = result ?: EmptyResult(), + ), + ) + } catch (e: CancellationException) { + // Intentional exception to the always-rethrow-CE convention, only at this boundary. + if (handlerJob == null) throw e // serial phase: inline dispatch is not response-suppressed + if (handlerJob.isCancelled) throw e // genuine peer or close cancel: suppress the response + // CE escaped a live handler job (e.g. a leaked inner withTimeout). Answer + // INTERNAL_ERROR so the peer does not hang until its own timeout. + respondWithError(capturedTransport, request, e) + } catch (cause: Throwable) { + if (handlerJob?.isCancelled == true) { + logger.trace { "Suppressing error response for request cancelled by peer (id: ${request.id})" } + return@withContext + } + respondWithError(capturedTransport, request, cause) + } + } + } + + private fun registerInFlight(connection: Connection, id: RequestId, job: Job) { + val previous = connection.inFlightRequestJobs.getAndUpdate { current -> current.putting(id, job) }[id] + if (previous != null) { + logger.warn { "Duplicate in-flight request id $id; replacing the previous handler job" } + } + job.invokeOnCompletion { + // Identity guard: never evict a successor job registered under a reused id. + connection.inFlightRequestJobs.update { current -> + if (current[id] === job) current.removing(id) else current + } + } + } + + private suspend fun respondWithError(transport: Transport, request: JSONRPCRequest, cause: Throwable) { + logger.error(cause) { "Error handling request: ${request.method} (id: ${request.id})" } + try { + val rpcError = if (cause is McpException) { + RPCError(code = cause.code, message = cause.message.orEmpty(), data = cause.data) + } else { + RPCError(code = RPCError.ErrorCode.INTERNAL_ERROR, message = cause.message ?: "Internal error") + } + transport.send(JSONRPCError(id = request.id, error = rpcError)) } catch (e: CancellationException) { throw e - } catch (cause: Throwable) { - logger.error(cause) { "Error handling request: ${request.method} (id: ${request.id})" } - - try { - val rpcError = if (cause is McpException) { - RPCError(code = cause.code, message = cause.message.orEmpty(), data = cause.data) - } else { - RPCError(code = RPCError.ErrorCode.INTERNAL_ERROR, message = cause.message ?: "Internal error") - } - transport?.send(JSONRPCError(id = request.id, error = rpcError)) - } catch (e: CancellationException) { - throw e - } catch (sendError: Throwable) { - logger.error(sendError) { - "Failed to send error response for request: ${request.method} (id: ${request.id})" - } - // Optionally implement fallback behavior here + } catch (sendError: Throwable) { + logger.error(sendError) { + "Failed to send error response for request: ${request.method} (id: ${request.id})" } + onError(sendError) } } @@ -354,6 +681,10 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio val handler = _progressHandlers.value[progressToken] if (handler == null) { + if (isRecentlyCancelled(progressToken)) { + logger.trace { "Ignoring progress for a locally cancelled request: $progressToken" } + return + } val error = Error( "Received a progress notification for an unknown token: ${McpJson.encodeToString(notification)}", ) @@ -370,7 +701,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio val oldResponseHandlers = _responseHandlers.getAndUpdate { current -> if (messageId != null && messageId in current) { - current.remove(messageId) + current.removing(messageId) } else { current } @@ -379,8 +710,12 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio val handler = oldResponseHandlers[messageId] if (handler != null) { - messageId?.let { msg -> _progressHandlers.update { it.remove(msg) } } + messageId?.let { msg -> _progressHandlers.update { it.removing(msg) } } } else { + if (isRecentlyCancelled(messageId)) { + logger.trace { "Ignoring response for a locally cancelled request: $messageId" } + return + } onError( IllegalStateException( "Received a response for an unknown message ID: ${McpJson.encodeToString(error ?: response)}", @@ -438,7 +773,8 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio public suspend fun request(request: Request, options: RequestOptions? = null): T { logger.trace { "Sending request: ${request.method}" } val result = CompletableDeferred() - val transport = transport ?: error("Not connected") + val connection = connectionRef.value ?: error("Not connected") + val transport = connection.transport if (this@Protocol.options?.enforceStrictCapabilities == true) { assertCapabilityForMethod(request.method) @@ -448,7 +784,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio options?.onProgress?.let { progressHandler -> logger.trace { "Registering progress handler for request id: $id" } _progressHandlers.update { current -> - current.put(id, progressHandler) + current.putting(id, progressHandler) } val paramsObject = (this.params as? JsonObject) ?: JsonObject(emptyMap()) @@ -467,10 +803,10 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio val jsonRpcRequestId = jsonRpcRequest.id _responseHandlers.update { current -> - current.put(jsonRpcRequestId) { response, error -> + current.putting(jsonRpcRequestId) { response, error -> if (error != null) { result.completeExceptionally(error) - return@put + return@putting } try { @@ -482,42 +818,84 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio } } - val cancel: suspend (Throwable) -> Unit = { reason: Throwable -> - _responseHandlers.update { current -> current.remove(jsonRpcRequestId) } - _progressHandlers.update { current -> current.remove(jsonRpcRequestId) } + if (connectionRef.value !== connection) { + // Connection closed or replaced during registration; nothing would complete this handler. + _responseHandlers.update { it.removing(jsonRpcRequestId) } + _progressHandlers.update { it.removing(jsonRpcRequestId) } + throw McpException(RPCError.ErrorCode.CONNECTION_CLOSED, "Connection closed") + } - val notification = CancelledNotification( - params = CancelledNotificationParams( - requestId = jsonRpcRequestId, - reason = reason.message ?: "Unknown", - ), - ) + val cancelPending: suspend (reason: Throwable, notifyPeer: Boolean) -> Unit = { reason, notifyPeer -> + // Remember before removing: a racing onResponse/onProgress that sees the handler gone + // must also see the id, or it reports a late message as unknown. Order is load-bearing. + rememberCancelledRequestId(jsonRpcRequestId) + _responseHandlers.update { current -> current.removing(jsonRpcRequestId) } + _progressHandlers.update { current -> current.removing(jsonRpcRequestId) } - val jsonRpcNotification = notification.toJSON() - - transport.send(jsonRpcNotification, options) + if (notifyPeer) { + val notification = CancelledNotification( + params = CancelledNotificationParams( + requestId = jsonRpcRequestId, + reason = reason.message ?: "Unknown", + ), + ) + transport.send(notification.toJSON(), options) + } result.completeExceptionally(reason) } + // The MCP spec forbids cancelling `initialize`; local cleanup still runs, + // but no notifications/cancelled goes on the wire (timeout and cancel paths alike). + val notifyPeerOnCancel = request.method != Method.Defined.Initialize + val timeout = options?.timeout ?: DEFAULT_REQUEST_TIMEOUT try { - withTimeout(timeout) { + val response = withTimeoutOrNull(timeout) { logger.trace { "Sending request message with id: $jsonRpcRequestId" } - this@Protocol.transport?.send(jsonRpcRequest, options) + transport.send(jsonRpcRequest, options) + result.await() } - return result.await() - } catch (cause: TimeoutCancellationException) { - logger.error { "Request timed out after ${timeout.inWholeMilliseconds}ms: ${request.method}" } - cancel( - McpException( + if (response == null) { + // Our own request timeout expired. An outer withTimeout's cancellation propagates + // through withTimeoutOrNull instead of returning null and is handled by the catch below. + logger.error { "Request timed out after ${timeout.inWholeMilliseconds}ms: ${request.method}" } + val timeoutError = McpException( code = RPCError.ErrorCode.REQUEST_TIMEOUT, message = "Request timed out", data = JsonObject(mutableMapOf("timeout" to JsonPrimitive(timeout.inWholeMilliseconds))), - ), - ) - result.cancel(cause) + ) + withContext(NonCancellable) { + try { + cancelPending(timeoutError, notifyPeerOnCancel) + } catch (e: Throwable) { + logger.warn(e) { "Failed to notify peer about timed-out request" } + onError(e) + } + } + result.cancel() + throw timeoutError + } + return response + } catch (cause: CancellationException) { + // Plain caller cancellation, or an outer withTimeout that propagated through + // withTimeoutOrNull (only our own deadline returns null there). NonCancellable so the + // cancellation notification still goes out even though the caller's job is cancelled. + // A secondary send failure must never replace the original CancellationException. + withContext(NonCancellable) { + try { + cancelPending(cause, notifyPeerOnCancel) + } catch (e: Throwable) { + logger.warn(e) { "Failed to notify peer about cancelled request" } + onError(e) + } + } throw cause + } finally { + // Backstop for the failed-send path. The success, timeout and cancellation paths already + // removed these handlers, so this only bites when the outbound send failed. + _responseHandlers.update { it.removing(jsonRpcRequestId) } + _progressHandlers.update { it.removing(jsonRpcRequestId) } } } @@ -538,6 +916,11 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio * Registers a handler to invoke when this protocol object receives a request with the given method. * * Note that this will replace any previous request handler for the same method. + * + * Handlers may run concurrently after initialization and are cancelled cooperatively, so always + * re-throw [CancellationException]. Replacing a built-in control-method handler (`ping`, + * `notifications/cancelled`, `notifications/progress`, `notifications/initialized`) with slow + * code degrades connection liveness. */ public inline fun setRequestHandler( method: Method, @@ -556,7 +939,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio val wrapped = wrapRequestHandler(method, block) _requestHandlers.update { current -> - current.put(method.value) { jSONRPCRequest, extraHandler -> + current.putting(method.value) { jSONRPCRequest, extraHandler -> val request = jSONRPCRequest.fromJSON() val response = wrapped(request as T, extraHandler) response @@ -582,17 +965,23 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio * Removes the request handler for the given method. */ public fun removeRequestHandler(method: Method) { - _requestHandlers.update { current -> current.remove(method.value) } + _requestHandlers.update { current -> current.removing(method.value) } } /** * Registers a handler to invoke when this protocol object receives a notification with the given method. * * Note that this will replace any previous notification handler for the same method. + * + * Handlers may run concurrently after initialization and are cancelled cooperatively, so always + * re-throw [CancellationException]. Replacing a built-in control-method handler (`ping`, + * `notifications/cancelled`, `notifications/progress`, `notifications/initialized`) with slow + * code degrades connection liveness. Replacing `notifications/cancelled` in particular disables + * inbound request cancellation entirely, since that handler is what cancels the running job. */ public fun setNotificationHandler(method: Method, handler: (notification: T) -> Deferred) { _notificationHandlers.update { current -> - current.put(method.value) { + current.putting(method.value) { @Suppress("UNCHECKED_CAST") handler(it.fromJSON() as T) } @@ -603,6 +992,22 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio * Removes the notification handler for the given method. */ public fun removeNotificationHandler(method: Method) { - _notificationHandlers.update { current -> current.remove(method.value) } + _notificationHandlers.update { current -> current.removing(method.value) } + } + + private class Connection( + val transport: Transport, + val handlerScope: CoroutineScope, + val executionSemaphore: Semaphore, + val maxInFlightHandlers: Int, + ) { + /** Init gate: inbound dispatch is inline/serial until this flips. */ + val concurrentDispatchEnabled = atomic(false) + + /** Admission tier: launched-but-not-completed handler jobs (running + parked). */ + val inFlightCount = atomic(0) + + /** In-flight request handler jobs, for `notifications/cancelled`. */ + val inFlightRequestJobs: AtomicRef> = atomic(persistentMapOf()) } } diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/methods.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/methods.kt index 19f4da5..90403e8 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/methods.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/methods.kt @@ -53,4 +53,13 @@ public sealed interface Method { */ @Serializable public data class Custom(override val value: String) : Method + + public companion object { + private val definedByValue: Map by lazy { + Defined.entries.associateBy { it.value } + } + + /** Resolves a wire method string to a [Defined] entry when known, [Custom] otherwise. */ + internal fun from(value: String): Method = definedByValue[value] ?: Custom(value) + } } 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 ab85e6d..6eab14a 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 @@ -59,18 +59,11 @@ internal object MethodSerializer : KSerializer { override val descriptor: SerialDescriptor = PrimitiveSerialDescriptor("io.modelcontextprotocol.kotlin.sdk.types.Method", PrimitiveKind.STRING) - private val definedMethods: Map by lazy { - Method.Defined.entries.associateBy { it.value } - } - override fun serialize(encoder: Encoder, value: Method) { encoder.encodeString(value.value) } - override fun deserialize(decoder: Decoder): Method { - val decodedString = decoder.decodeString() - return definedMethods[decodedString] ?: Method.Custom(decodedString) - } + override fun deserialize(decoder: Decoder): Method = Method.from(decoder.decodeString()) } // ============================================================================ diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyTest.kt new file mode 100644 index 0000000..b3c95ac --- /dev/null +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyTest.kt @@ -0,0 +1,524 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.kotest.assertions.throwables.shouldThrow +import io.kotest.matchers.collections.shouldHaveSize +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldContain +import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult +import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse +import io.modelcontextprotocol.kotlin.sdk.types.McpException +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.PingRequest +import io.modelcontextprotocol.kotlin.sdk.types.RPCError +import io.modelcontextprotocol.kotlin.sdk.types.RequestId +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.async +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.StandardTestDispatcher +import kotlinx.coroutines.test.TestScope +import kotlinx.coroutines.test.advanceTimeBy +import kotlinx.coroutines.test.advanceUntilIdle +import kotlinx.coroutines.test.runCurrent +import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.int +import kotlinx.serialization.json.jsonPrimitive +import kotlinx.serialization.json.put +import kotlin.coroutines.cancellation.CancellationException +import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds + +class ProtocolConcurrencyTest { + + private fun TestScope.newProtocol(): TestProtocol = + TestProtocol(ProtocolOptions(handlerCoroutineContext = StandardTestDispatcher(testScheduler))) + + @Test + fun `a suspended slow handler does not block a later fast request`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/slow") delay(10.seconds) + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow")) + transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast")) + runCurrent() + + responsesOn(transport).map { it.id } shouldBe listOf(RequestId(2L)) + + advanceTimeBy(11.seconds) + runCurrent() + responsesOn(transport).map { it.id } shouldBe listOf(RequestId(2L), RequestId(1L)) + } + + @Test + fun `dispatch is serial until concurrent dispatch is enabled`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + // gate NOT flipped + + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/slow") delay(10.seconds) + EmptyResult() + } + + val delivery = launch { + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow")) + transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast")) + } + runCurrent() + // inline handling: the delivering coroutine is parked inside the slow handler + delivery.isCompleted shouldBe false + responsesOn(transport) shouldBe emptyList() + + advanceTimeBy(11.seconds) + runCurrent() + delivery.isCompleted shouldBe true + responsesOn(transport).map { it.id } shouldBe listOf(RequestId(1L), RequestId(2L)) + } + + // the hook fires before handler lookup and cannot be disabled by a user handler + @Test + fun `initialized notification invokes the router hook even when a user handler is registered`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + var userHandlerCalled = false + protocol.setNotificationHandler(Method.Defined.NotificationsInitialized) { + userHandlerCalled = true + COMPLETED + } + transport.deliver(JSONRPCNotification(method = Method.Defined.NotificationsInitialized.value)) + runCurrent() + + protocol.initializedNotificationCount shouldBe 1 + userHandlerCalled shouldBe true + } + + @Test + fun `cancelled notification cancels the in-flight handler and suppresses the response`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val entered = CompletableDeferred() + val sawCancellation = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + entered.complete(Unit) + try { + delay(10.seconds) + } catch (e: CancellationException) { + sawCancellation.complete(Unit) + throw e + } + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(5L), method = "test/slow")) + runCurrent() + entered.await() + + transport.deliver(cancelledNotification(requestId = RequestId(5L), reason = "user gave up")) + runCurrent() + sawCancellation.await() + advanceUntilIdle() + + responsesOn(transport) shouldBe emptyList() + errorsOn(transport) shouldBe emptyList() // JSONRPCError messages on the wire + protocol.errors shouldBe emptyList() + } + + @Test + fun `cancelled notification for an unknown id is a silent no-op`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + transport.deliver(cancelledNotification(requestId = RequestId(404L), reason = "nope")) + advanceUntilIdle() + + protocol.errors shouldBe emptyList() + } + + @Test + fun `handler leaking CancellationException without being cancelled produces INTERNAL_ERROR`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + protocol.fallbackRequestHandler = { _, _ -> throw CancellationException("leaked") } + + transport.deliver(JSONRPCRequest(id = RequestId(6L), method = "test/leak")) + advanceUntilIdle() + + val error = errorsOn(transport).single() + error.id shouldBe RequestId(6L) + error.error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR + } + + // suppression race variants + @Test + fun `handler that swallows cancellation and returns normally still gets no response`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val entered = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + entered.complete(Unit) + try { + delay(10.seconds) + } catch (_: CancellationException) { + // swallow (anti-pattern) and return normally + } + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(7L), method = "test/swallow")) + runCurrent() + entered.await() + transport.deliver(cancelledNotification(requestId = RequestId(7L), reason = "stop")) + advanceUntilIdle() + + responsesOn(transport) shouldBe emptyList() + } + + @Test + fun `handler throwing non-CE while cancelled gets no error response`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val entered = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + entered.complete(Unit) + try { + delay(10.seconds) + } catch (_: CancellationException) { + throw IllegalStateException("boom") + } + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(8L), method = "test/throw")) + runCurrent() + entered.await() + transport.deliver(cancelledNotification(requestId = RequestId(8L), reason = "stop")) + advanceUntilIdle() + + errorsOn(transport) shouldBe emptyList() + protocol.errors shouldBe emptyList() + } + + @Test + fun `suspension-free notification handlers observe strict arrival order`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val seen = mutableListOf() + protocol.fallbackNotificationHandler = { notification -> + seen.add((notification.params as JsonObject).getValue("i").jsonPrimitive.int) + } + + repeat(20) { i -> + transport.deliver( + JSONRPCNotification(method = "test/event", params = buildJsonObject { put("i", i) }), + ) + } + advanceUntilIdle() + + seen shouldBe (0 until 20).toList() + } + + @Test + fun `reconnect resets the dispatch gate to serial`() = runTest { + val protocol = newProtocol() + val transport1 = RecordingTransport() + protocol.connect(transport1) + protocol.enableConcurrency() + protocol.close() + + val transport2 = RecordingTransport() + protocol.connect(transport2) + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/slow") delay(10.seconds) + EmptyResult() + } + + val delivery = launch { + transport2.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow")) + transport2.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast")) + } + runCurrent() + delivery.isCompleted shouldBe false // serial again: gate reset + advanceTimeBy(11.seconds) + runCurrent() + responsesOn(transport2).map { it.id } shouldBe listOf(RequestId(1L), RequestId(2L)) + } + + @Test + fun `close cancels in-flight handlers suppresses responses and fails pending outbound requests`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val entered = CompletableDeferred() + val sawCancellation = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + entered.complete(Unit) + try { + delay(10.seconds) + } catch (e: CancellationException) { + sawCancellation.complete(Unit) + throw e + } + EmptyResult() + } + transport.deliver(JSONRPCRequest(id = RequestId(9L), method = "test/slow")) + runCurrent() + entered.await() + + // shouldThrow inside the async so its failure does not cancel the (non-supervisor) test scope. + val pending = async { shouldThrow { protocol.request(PingRequest()) } } + transport.awaitRequest() // outbound ping on the wire + + protocol.close() + runCurrent() + + sawCancellation.await() + responsesOn(transport) shouldBe emptyList() + val thrown = pending.await() + thrown.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED + } + + @Test + fun `message delivered after close is dropped without crash or response`() = runTest { + val protocol = newProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + protocol.close() + + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/anything")) + advanceUntilIdle() + + responsesOn(transport) shouldBe emptyList() + errorsOn(transport) shouldBe emptyList() + } + + @Test + fun `at most maxConcurrentHandlers handlers run concurrently while control messages bypass`() = runTest { + val protocol = TestProtocol( + ProtocolOptions( + handlerCoroutineContext = StandardTestDispatcher(testScheduler), + maxConcurrentHandlers = 2, + ), + ) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + var entered = 0 + val release = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + entered += 1 + release.await() + EmptyResult() + } + + repeat(4) { i -> transport.deliver(JSONRPCRequest(id = RequestId(i.toLong()), method = "test/slow")) } + runCurrent() + entered shouldBe 2 // exactly two running, two parked on the semaphore + + // ping (bypass) is answered while saturated + transport.deliver(JSONRPCRequest(id = RequestId(100L), method = Method.Defined.Ping.value)) + runCurrent() + responsesOn(transport).map { it.id } shouldBe listOf(RequestId(100L)) + + // responses (bypass, inline) are processed while saturated + val pending = async { protocol.request(PingRequest()) } + runCurrent() + val outbound = transport.awaitRequest() + transport.deliver(JSONRPCResponse(id = outbound.id, result = EmptyResult())) + runCurrent() + pending.isCompleted shouldBe true + + release.complete(Unit) + advanceUntilIdle() + entered shouldBe 4 + responsesOn(transport) shouldHaveSize 5 + } + + @Test + fun `flood beyond maxInFlightHandlers is rejected fail-fast and the read loop never suspends`() = runTest { + val protocol = TestProtocol( + ProtocolOptions( + handlerCoroutineContext = StandardTestDispatcher(testScheduler), + maxConcurrentHandlers = 1, + maxInFlightHandlers = 4, + ), + ) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val release = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + release.await() + EmptyResult() + } + + // 10 requests delivered back-to-back; delivery must complete without advancing + // virtual time — i.e. the read loop was never suspended on admission. + val delivery = launch { + repeat(10) { i -> transport.deliver(JSONRPCRequest(id = RequestId(i.toLong()), method = "test/slow")) } + } + runCurrent() + delivery.isCompleted shouldBe true + + // 4 admitted (1 running + 3 parked), 6 rejected immediately + val busyErrors = errorsOn(transport) + busyErrors shouldHaveSize 6 + busyErrors.forEach { + it.error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR + it.error.message shouldBe "Server is busy: too many in-flight messages" + } + + // ping still answered at full saturation + transport.deliver(JSONRPCRequest(id = RequestId(100L), method = Method.Defined.Ping.value)) + runCurrent() + responsesOn(transport).map { it.id } shouldBe listOf(RequestId(100L)) + + release.complete(Unit) + advanceUntilIdle() + responsesOn(transport) shouldHaveSize 5 // ping + the 4 admitted + } + + @Test + fun `overflowing notifications are dropped and reported via onError`() = runTest { + val protocol = TestProtocol( + ProtocolOptions( + handlerCoroutineContext = StandardTestDispatcher(testScheduler), + maxConcurrentHandlers = 1, + maxInFlightHandlers = 2, + ), + ) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + // parks one handler with a request: holds the sole execution permit and one in-flight slot + val release = CompletableDeferred() + protocol.fallbackRequestHandler = { _, _ -> + release.await() + EmptyResult() + } + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/hog")) + runCurrent() + + // notification 1 is admitted (parked on the semaphore, behind the hogged permit); + // notification 2 overflows the admission cap and is dropped instead of parked. + val notificationRelease = CompletableDeferred() + protocol.fallbackNotificationHandler = { notificationRelease.await() } + transport.deliver(JSONRPCNotification(method = "test/event")) + transport.deliver(JSONRPCNotification(method = "test/event")) + runCurrent() + + protocol.errors shouldHaveSize 1 + protocol.errors.single().message.orEmpty() shouldContain "too many in-flight messages" + errorsOn(transport) shouldBe emptyList() // no wire error for dropped notifications + + release.complete(Unit) + notificationRelease.complete(Unit) + advanceUntilIdle() + } + + @Test + fun `cancelling saturating handlers releases permits and lets parked handlers run`() = runTest { + val protocol = TestProtocol( + ProtocolOptions( + handlerCoroutineContext = StandardTestDispatcher(testScheduler), + maxConcurrentHandlers = 1, + ), + ) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val secondCompleted = CompletableDeferred() + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/hog") { + delay(1000.seconds) // cancellable park + error("unreachable") + } + secondCompleted.complete(Unit) + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/hog")) // holds the permit + transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/next")) // parked on the semaphore + runCurrent() + secondCompleted.isCompleted shouldBe false + + // cancel the hog THROUGH the bypass while saturated + transport.deliver(cancelledNotification(requestId = RequestId(1L), reason = "make room")) + advanceUntilIdle() + + secondCompleted.isCompleted shouldBe true + responsesOn(transport).map { it.id } shouldBe listOf(RequestId(2L)) // hog suppressed, next answered + } + + @Test + fun `a parked handler can itself be cancelled before it ever runs`() = runTest { + val protocol = TestProtocol( + ProtocolOptions( + handlerCoroutineContext = StandardTestDispatcher(testScheduler), + maxConcurrentHandlers = 1, + ), + ) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/hog") { + delay(1000.seconds) // cancellable park + error("unreachable") + } + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/hog")) // holds the permit + transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/next")) // parked on the semaphore + runCurrent() + + // cancel the parked handler before it ever acquires the permit, then cancel the hog too + transport.deliver(cancelledNotification(requestId = RequestId(2L), reason = "give up while parked")) + transport.deliver(cancelledNotification(requestId = RequestId(1L), reason = "make room")) + advanceUntilIdle() + + // no response for id=2 at all: its withPermit acquire was cancelled before onRequest ever ran + responsesOn(transport) shouldBe emptyList() + errorsOn(transport) shouldBe emptyList() + } +} diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptionsTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptionsTest.kt new file mode 100644 index 0000000..81bf24a --- /dev/null +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptionsTest.kt @@ -0,0 +1,16 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.kotest.matchers.shouldBe +import kotlinx.coroutines.Dispatchers +import kotlin.test.Test + +class ProtocolOptionsTest { + + @Test + fun `defaults expose concurrency knobs`() { + val options = ProtocolOptions() + options.handlerCoroutineContext shouldBe Dispatchers.Default + options.maxConcurrentHandlers shouldBe 64 + options.maxInFlightHandlers shouldBe 256 + } +} diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOutboundCancellationTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOutboundCancellationTest.kt new file mode 100644 index 0000000..38ad556 --- /dev/null +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOutboundCancellationTest.kt @@ -0,0 +1,167 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.kotest.assertions.throwables.shouldThrow +import io.kotest.matchers.collections.shouldHaveSize +import io.kotest.matchers.nulls.shouldNotBeNull +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldContain +import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams +import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest +import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse +import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION +import io.modelcontextprotocol.kotlin.sdk.types.McpException +import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.PingRequest +import io.modelcontextprotocol.kotlin.sdk.types.RPCError +import io.modelcontextprotocol.kotlin.sdk.types.RequestId +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.TimeoutCancellationException +import kotlinx.coroutines.async +import kotlinx.coroutines.cancelAndJoin +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.advanceTimeBy +import kotlinx.coroutines.test.runCurrent +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withTimeout +import kotlinx.serialization.json.decodeFromJsonElement +import kotlinx.serialization.json.jsonObject +import kotlinx.serialization.json.jsonPrimitive +import kotlinx.serialization.json.long +import kotlin.coroutines.cancellation.CancellationException +import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds + +class ProtocolOutboundCancellationTest { + + @Test + fun `request times out awaiting the response and sends CancelledNotification with timeout reason`() = runTest { + val protocol = TestProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + val inFlight = async { + shouldThrow { + protocol.request(PingRequest(), RequestOptions(timeout = 5.seconds)) + } + } + val sent = transport.awaitRequest() // request is on the wire; peer never responds + advanceTimeBy(6.seconds) + runCurrent() + + val thrown = inFlight.await() + thrown.code shouldBe RPCError.ErrorCode.REQUEST_TIMEOUT + thrown.data?.jsonObject?.get("timeout")?.jsonPrimitive?.long shouldBe 5.seconds.inWholeMilliseconds + + val cancelledJson = transport.sentWithOptions + .map { it.first } + .filterIsInstance() + .single { it.method == Method.Defined.NotificationsCancelled.value } + val cancelledParams = McpJson.decodeFromJsonElement(cancelledJson.params!!) + cancelledParams.requestId shouldBe sent.id + cancelledParams.reason.shouldNotBeNull() shouldContain "timed out" + } + + @Test + fun `outer withTimeout around request propagates the original timeout exception and notifies the peer`() = runTest { + val protocol = TestProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + val inFlight = async { + shouldThrow { + withTimeout(1.seconds) { + protocol.request(PingRequest()) + } + } + } + transport.awaitRequest() // request is on the wire; peer never responds + advanceTimeBy(2.seconds) + runCurrent() + inFlight.await() + + transport.sentWithOptions.map { it.first }.filterIsInstance() + .filter { it.method == Method.Defined.NotificationsCancelled.value } shouldHaveSize 1 + } + + @Test + fun `caller cancellation sends CancelledNotification and rethrows the original exception`() = runTest { + val protocol = TestProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + val cancelled = CompletableDeferred() + val job = launch { + try { + protocol.request(PingRequest()) + } catch (e: CancellationException) { + cancelled.complete(e) + throw e + } + } + transport.awaitRequest() + job.cancelAndJoin() + + cancelled.await() + val wire = transport.sentWithOptions.map { it.first }.filterIsInstance() + .filter { it.method == Method.Defined.NotificationsCancelled.value } + wire shouldHaveSize 1 + } + + @Test + fun `cancelling the initialize request performs local cleanup but sends no CancelledNotification`() = runTest { + val protocol = TestProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + val job = launch { + protocol.request( + InitializeRequest( + InitializeRequestParams( + protocolVersion = LATEST_PROTOCOL_VERSION, + capabilities = ClientCapabilities(), + clientInfo = Implementation(name = "t", version = "1"), + ), + ), + ) + } + val sent = transport.awaitRequest() + sent.method shouldBe Method.Defined.Initialize.value + job.cancelAndJoin() + + transport.sentWithOptions.map { it.first }.filterIsInstance() + .filter { it.method == Method.Defined.NotificationsCancelled.value } shouldHaveSize 0 + protocol.responseHandlers shouldBe emptyMap() + } + + @Test + fun `late response for a cancelled request is ignored silently`() = runTest { + val protocol = TestProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + val job = launch { protocol.request(PingRequest()) } + val sent = transport.awaitRequest() + job.cancelAndJoin() + + transport.deliver(JSONRPCResponse(id = sent.id, result = EmptyResult())) + + protocol.errors shouldHaveSize 0 // no "unknown message ID" onError + } + + @Test + fun `response for a genuinely unknown id still reports onError`() = runTest { + val protocol = TestProtocol() + val transport = RecordingTransport() + protocol.connect(transport) + + transport.deliver(JSONRPCResponse(id = RequestId(999L), result = EmptyResult())) + + protocol.errors shouldHaveSize 1 + } +} 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 b726c2f..2846ccc 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 @@ -5,6 +5,9 @@ import io.kotest.matchers.collections.shouldContainExactly import io.kotest.matchers.collections.shouldHaveSize import io.kotest.matchers.nulls.shouldNotBeNull 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.CustomRequest import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage @@ -13,11 +16,16 @@ import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse import io.modelcontextprotocol.kotlin.sdk.types.McpJson import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.PingRequest +import io.modelcontextprotocol.kotlin.sdk.types.ProgressNotification +import io.modelcontextprotocol.kotlin.sdk.types.ProgressNotificationParams +import io.modelcontextprotocol.kotlin.sdk.types.ProgressToken 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 kotlinx.coroutines.async -import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.launch import kotlinx.coroutines.test.runTest import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonObjectBuilder @@ -30,6 +38,7 @@ import kotlinx.serialization.json.jsonPrimitive import kotlin.coroutines.cancellation.CancellationException import kotlin.test.BeforeTest import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds class ProtocolTest { private lateinit var protocol: TestProtocol @@ -185,64 +194,174 @@ class ProtocolTest { transport.deliver(JSONRPCResponse(sent.id, EmptyResult())) inFlight.await() } -} -private class TestProtocol : Protocol(null) { - val errors = mutableListOf() + @Test + fun `request handler receives enriched extra and ambient context element`() = runTest { + protocol.connect(transport) - override fun onError(error: Throwable) { - errors.add(error) + var seenByParameter: RequestHandlerExtra? = null + var seenAmbient: RequestHandlerExtra? = null + protocol.fallbackRequestHandler = { _, extra -> + seenByParameter = extra + seenAmbient = currentRequestHandlerExtra() + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(7L), method = "custom/echo")) + + seenByParameter shouldNotBe null + seenByParameter?.requestId shouldBe RequestId(7L) + seenByParameter?.method shouldBe Method.Custom("custom/echo") + // the SAME instance flows through both delivery paths + seenAmbient shouldBe seenByParameter } - override fun assertCapabilityForMethod(method: Method) { - // noop - } - override fun assertNotificationCapability(method: Method) { - // noop - } - override fun assertRequestHandlerCapability(method: Method) { - // noop - } -} + @Test + fun `extra method resolves defined methods to Defined entries`() = runTest { + protocol.connect(transport) -private class RecordingTransport : Transport { - private val sentMessages = Channel(Channel.UNLIMITED) - private var onMessageCallback: (suspend (JSONRPCMessage) -> Unit)? = null - private var onCloseCallback: (() -> Unit)? = null + var seenMethod: Method? = null + protocol.setRequestHandler(Method.Defined.Ping) { _, extra -> + seenMethod = extra.method + EmptyResult() + } + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "ping")) - override suspend fun start() { - // noop + seenMethod shouldBe Method.Defined.Ping } - override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) { - sentMessages.send(message) + @Test + fun `extra sendNotification stamps relatedRequestId`() = runTest { + protocol.connect(transport) + + protocol.fallbackRequestHandler = { _, extra -> + extra.sendNotification( + ProgressNotification( + ProgressNotificationParams(progressToken = ProgressToken(1L), progress = 0.5), + ), + ) + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(42L), method = "custom/progressing")) + + val sent = transport.sentWithOptions.first { it.first is JSONRPCNotification } + sent.second?.relatedRequestId shouldBe RequestId(42L) } - override suspend fun close() { - onCloseCallback?.invoke() + @Test + fun `extra sendRequest stamps relatedRequestId and preserves caller options`() = runTest { + protocol.connect(transport) + + protocol.fallbackRequestHandler = { _, extra -> + extra.sendRequest( + PingRequest(), + RequestOptions(resumptionToken = "tok", onProgress = { }, timeout = 30.seconds), + ) + EmptyResult() + } + + // The serial phase runs the handler inline inside deliver(), and the handler suspends + // awaiting the nested request's response — drive delivery from a child coroutine. + val delivery = launch { + transport.deliver(JSONRPCRequest(id = RequestId(42L), method = "custom/nested")) + } + + val outbound = transport.awaitRequest() + val recorded = transport.sentWithOptions + .first { (it.first as? JSONRPCRequest)?.id == outbound.id } + .second + val options = recorded.shouldBeInstanceOf() + options.relatedRequestId shouldBe RequestId(42L) + options.resumptionToken shouldBe "tok" + options.timeout shouldBe 30.seconds + options.onProgress shouldNotBe null + + transport.deliver(JSONRPCResponse(id = outbound.id, result = EmptyResult())) + delivery.join() } - override fun onClose(block: () -> Unit) { - onCloseCallback = block + @Test + fun `currentRequestHandlerExtra returns null outside handlers`() = runTest { + currentRequestHandlerExtra() shouldBe null } - override fun onError(block: (Throwable) -> Unit) { - // noop + @Test + fun `connect while already connected throws IllegalStateException`() = runTest { + protocol.connect(transport) + shouldThrow { + protocol.connect(RecordingTransport()) + } } - override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) { - onMessageCallback = block + @Test + fun `stale onClose from a previous transport does not tear down the successor connection`() = runTest { + protocol.connect(transport) + val staleCloseCallback = transport.closeCallback ?: error("onClose not registered") + + protocol.close() // disconnect first transport (fires its own doClose) + val second = RecordingTransport() + protocol.connect(second) // reconnect + + staleCloseCallback() // late duplicate close signal from transport #1 + + // successor connection must still be alive + protocol.transport shouldBe second } - suspend fun awaitRequest(): JSONRPCRequest { - val message = sentMessages.receive() - return message as? JSONRPCRequest - ?: error("Expected JSONRPCRequest but received ${message::class.simpleName}") + @Test + fun `failed transport start rolls back the connection and allows reconnect`() = runTest { + val failing = object : Transport { + override suspend fun start(): Unit = error("boom") + override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {} + override suspend fun close() {} + override fun onClose(block: () -> Unit) {} + override fun onError(block: (Throwable) -> Unit) {} + override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {} + } + shouldThrow { protocol.connect(failing) } + protocol.transport shouldBe null // rolled back + protocol.connect(transport) // reconnect succeeds, no "already connected" + protocol.transport shouldBe transport } - suspend fun deliver(message: JSONRPCMessage) { - val callback = onMessageCallback ?: error("onMessage callback not registered") - callback(message) + @Test + fun `request cleans up its handlers when the transport send fails`() = runTest { + // Send fails without closing the connection, so the leak is not masked by doClose(). + val failingSend = object : Transport { + override suspend fun start() {} + override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?): Unit = + throw IllegalStateException("send failed") + override suspend fun close() {} + override fun onClose(block: () -> Unit) {} + override fun onError(block: (Throwable) -> Unit) {} + override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {} + } + protocol.connect(failingSend) + + shouldThrow { + protocol.request(PingRequest(), RequestOptions(onProgress = {})) + } + + protocol.responseHandlers.size shouldBe 0 + protocol.progressHandlers.size shouldBe 0 + } + + @Test + fun `connect rejects non-positive maxConcurrentHandlers`() = runTest { + val error = shouldThrow { + TestProtocol(ProtocolOptions(maxConcurrentHandlers = 0)).connect(RecordingTransport()) + } + error.message.shouldNotBeNull() shouldContain "maxConcurrentHandlers" + } + + @Test + fun `connect rejects maxInFlightHandlers below maxConcurrentHandlers`() = runTest { + val error = shouldThrow { + TestProtocol(ProtocolOptions(maxConcurrentHandlers = 4, maxInFlightHandlers = 0)) + .connect(RecordingTransport()) + } + error.message.shouldNotBeNull() shouldContain "maxInFlightHandlers" } } diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt new file mode 100644 index 0000000..9e8dd60 --- /dev/null +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolTestFakes.kt @@ -0,0 +1,112 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse +import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.RequestId +import kotlinx.coroutines.channels.Channel +import kotlinx.serialization.json.encodeToJsonElement + +/** + * Shared [Protocol] test double: records [onError] calls, counts `notifications/initialized` + * router-hook invocations, and no-ops all capability assertions. + * + * The concurrency gate is flipped only via [enableConcurrency]; nothing in this double flips it + * implicitly, so tests keep today's serial semantics unless they opt in. + */ +internal class TestProtocol(options: ProtocolOptions? = null) : Protocol(options) { + val errors = mutableListOf() + var initializedNotificationCount = 0 + + fun enableConcurrency() { + enableConcurrentDispatch() + } + + override fun onInitializedNotification() { + initializedNotificationCount++ + // deliberately does NOT flip the gate — tests control the flip explicitly + } + + override fun onError(error: Throwable) { + errors.add(error) + } + + override fun assertCapabilityForMethod(method: Method) { + // noop + } + override fun assertNotificationCapability(method: Method) { + // noop + } + public override fun assertRequestHandlerCapability(method: Method) { + // noop + } +} + +/** Shared [Transport] test double: records every sent message (with its options) and replays inbound ones. */ +internal class RecordingTransport : Transport { + val sentMessages = Channel(Channel.UNLIMITED) + private var onMessageCallback: (suspend (JSONRPCMessage) -> Unit)? = null + private var onCloseCallback: (() -> Unit)? = null + + val sentWithOptions = mutableListOf>() + + var closeCallback: (() -> Unit)? = null + private set + + override suspend fun start() { + // noop + } + + override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) { + sentWithOptions.add(message to options) + sentMessages.send(message) + } + + override suspend fun close() { + onCloseCallback?.invoke() + } + + override fun onClose(block: () -> Unit) { + closeCallback = block + onCloseCallback = block + } + + override fun onError(block: (Throwable) -> Unit) { + // noop + } + + override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) { + onMessageCallback = block + } + + suspend fun awaitRequest(): JSONRPCRequest { + while (true) { + val message = sentMessages.receive() + if (message is JSONRPCRequest) return message + } + } + + /** Cross-thread-safe reader used by real-dispatcher tests; [sentWithOptions] is single-thread only. */ + suspend fun awaitSent(): JSONRPCMessage = sentMessages.receive() + + suspend fun deliver(message: JSONRPCMessage) { + val callback = onMessageCallback ?: error("onMessage callback not registered") + callback(message) + } +} + +internal fun responsesOn(transport: RecordingTransport): List = + transport.sentWithOptions.map { it.first }.filterIsInstance() + +internal fun errorsOn(transport: RecordingTransport): List = + transport.sentWithOptions.map { it.first }.filterIsInstance() + +internal fun cancelledNotification(requestId: RequestId, reason: String): JSONRPCNotification = JSONRPCNotification( + method = Method.Defined.NotificationsCancelled.value, + params = McpJson.encodeToJsonElement(CancelledNotificationParams(requestId = requestId, reason = reason)), +) diff --git a/kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyJvmTest.kt b/kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyJvmTest.kt new file mode 100644 index 0000000..60658f8 --- /dev/null +++ b/kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyJvmTest.kt @@ -0,0 +1,114 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.kotest.matchers.shouldBe +import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse +import io.modelcontextprotocol.kotlin.sdk.types.RequestId +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeout +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.int +import kotlinx.serialization.json.jsonPrimitive +import kotlinx.serialization.json.put +import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds + +/** + * Real-dispatcher counterparts to [ProtocolConcurrencyTest]. + * + * These run under [runBlocking] (real time, real threads on [Dispatchers.Default]) rather than + * `runTest`: a virtual-time scheduler auto-advances past the `withTimeout` guards and cannot + * observe genuine cross-thread parallelism, so the concurrency assertions could pass vacuously. + */ +class ProtocolConcurrencyJvmTest { + + @Test + fun `fast request completes while a slow handler is parked on Dispatchers Default`() = runBlocking { + val protocol = TestProtocol(ProtocolOptions(handlerCoroutineContext = Dispatchers.Default)) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + val slowEntered = CompletableDeferred() + val releaseSlow = CompletableDeferred() + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/slow") { + slowEntered.complete(Unit) + releaseSlow.await() + } + EmptyResult() + } + + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow")) + withTimeout(5.seconds) { slowEntered.await() } + transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast")) + + val fast = withTimeout(5.seconds) { transport.awaitSent() } + (fast as JSONRPCResponse).id shouldBe RequestId(2L) + + releaseSlow.complete(Unit) + val slow = withTimeout(5.seconds) { transport.awaitSent() } + (slow as JSONRPCResponse).id shouldBe RequestId(1L) + Unit + } + + @Test + fun `dispatch stays serial before the gate flips on Dispatchers Default`() = runBlocking { + val protocol = TestProtocol(ProtocolOptions(handlerCoroutineContext = Dispatchers.Default)) + val transport = RecordingTransport() + protocol.connect(transport) + // gate NOT flipped + + val slowEntered = CompletableDeferred() + val releaseSlow = CompletableDeferred() + protocol.fallbackRequestHandler = { request, _ -> + if (request.method == "test/slow") { + slowEntered.complete(Unit) + releaseSlow.await() + } + EmptyResult() + } + + // A single delivering coroutine models the transport read loop: the second deliver cannot + // start until the first returns, so a parked slow handler blocks the fast request. + val delivery = launch { + transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow")) + transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast")) + } + withTimeout(5.seconds) { slowEntered.await() } + + // The delivering coroutine is parked inside the slow handler; fast was never delivered. + transport.sentMessages.tryReceive().isFailure shouldBe true + + releaseSlow.complete(Unit) + val first = withTimeout(5.seconds) { transport.awaitSent() } + (first as JSONRPCResponse).id shouldBe RequestId(1L) + val second = withTimeout(5.seconds) { transport.awaitSent() } + (second as JSONRPCResponse).id shouldBe RequestId(2L) + delivery.join() + } + + @Test + fun `suspension-free notification handlers keep arrival order on Dispatchers Default`() = runBlocking { + val protocol = TestProtocol(ProtocolOptions(handlerCoroutineContext = Dispatchers.Default)) + val transport = RecordingTransport() + protocol.connect(transport) + protocol.enableConcurrency() + + // safe: appended only from the delivering coroutine (UNDISPATCHED, suspension-free) + val seen = mutableListOf() + protocol.fallbackNotificationHandler = { notification -> + seen.add((notification.params as JsonObject).getValue("i").jsonPrimitive.int) + } + repeat(100) { i -> + transport.deliver(JSONRPCNotification(method = "test/event", params = buildJsonObject { put("i", i) })) + } + seen shouldBe (0 until 100).toList() + } +} diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index 6094b7c..51ac279 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -169,8 +169,8 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/ServerOptions : io/ public fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;)V public fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;Z)V public synthetic fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V - public fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;)V - public synthetic fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;Lkotlin/coroutines/CoroutineContext;)V + public synthetic fun (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;Lkotlin/coroutines/CoroutineContext;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun getCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities; public final fun getResourceTemplateMatcherFactory ()Lio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory; } @@ -198,6 +198,7 @@ public class io/modelcontextprotocol/kotlin/sdk/server/ServerSession : io/modelc public fun onClose ()V public final fun onClose (Lkotlin/jvm/functions/Function0;)V public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V + protected fun onInitializedNotification ()V public final fun ping (Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun sendElicitationComplete (Lio/modelcontextprotocol/kotlin/sdk/types/ElicitationCompleteNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun sendLoggingMessage (Lio/modelcontextprotocol/kotlin/sdk/types/LoggingMessageNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt index c0d3e84..51a2389 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt @@ -53,9 +53,11 @@ import io.modelcontextprotocol.kotlin.sdk.utils.ResourceTemplateMatcher import io.modelcontextprotocol.kotlin.sdk.utils.ResourceTemplateMatcherFactory import kotlinx.coroutines.CancellationException import kotlinx.coroutines.Deferred +import kotlinx.coroutines.Dispatchers import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.buildJsonObject import kotlinx.serialization.json.put +import kotlin.coroutines.CoroutineContext import kotlin.jvm.JvmOverloads import kotlin.time.ExperimentalTime @@ -69,12 +71,17 @@ private val logger = KotlinLogging.logger {} * Forwarded to [ProtocolOptions.enforceStrictCapabilities]. * @property resourceTemplateMatcherFactory The factory used to create [ResourceTemplateMatcher] instances * for matching resource URIs against registered templates. Defaults to [PathSegmentTemplateMatcher.factory]. + * @param handlerCoroutineContext Coroutine context for inbound handlers. See [ProtocolOptions.handlerCoroutineContext]. */ public class ServerOptions( public val capabilities: ServerCapabilities, enforceStrictCapabilities: Boolean = true, public val resourceTemplateMatcherFactory: ResourceTemplateMatcherFactory = PathSegmentTemplateMatcher.factory, -) : ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities) { + handlerCoroutineContext: CoroutineContext = Dispatchers.Default, +) : ProtocolOptions( + enforceStrictCapabilities = enforceStrictCapabilities, + handlerCoroutineContext = handlerCoroutineContext, +) { @JvmOverloads public constructor( capabilities: ServerCapabilities, diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt index d215bd5..880f1cf 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt @@ -108,10 +108,18 @@ public open class ServerSession( } } + override fun onInitializedNotification() { + // Enable concurrency before the user's onInitialized callback, which runs synchronously. + enableConcurrentDispatch() + } + internal val clientConnection: ClientConnection = ClientConnectionImpl(this) /** * Registers a callback to be invoked when the server has completed initialization. + * + * The callback must be synchronous and fast: it runs on the message-dispatch path for + * `notifications/initialized`, after concurrent dispatch has been enabled for the session. */ public fun onInitialized(block: () -> Unit) { val old = _onInitialized diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt index c213350..1f8ac39 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt @@ -15,10 +15,12 @@ import io.ktor.server.sse.ServerSSESession import io.ktor.util.collections.ConcurrentMap import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions +import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams import io.modelcontextprotocol.kotlin.sdk.types.DEFAULT_NEGOTIATED_PROTOCOL_VERSION import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCEmptyMessage import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse import io.modelcontextprotocol.kotlin.sdk.types.McpJson @@ -27,6 +29,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.RPCError import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode.REQUEST_TIMEOUT import io.modelcontextprotocol.kotlin.sdk.types.RequestId import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS +import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.NonCancellable import kotlinx.coroutines.awaitCancellation import kotlinx.coroutines.job @@ -175,6 +178,20 @@ public class StreamableHttpServerTransport(private val configuration: Configurat private val requestToStreamMapping: ConcurrentMap = ConcurrentMap() private val requestToResponseMapping: ConcurrentMap = ConcurrentMap() + /** + * JSON-response mode only: one entry per in-flight POST, completed by [send] with the collected + * responses once every request delivered by that POST has been answered. [handlePostRequest] + * awaits it so the HTTP exchange stays open while handlers run concurrently. + */ + private val pendingJsonResponses: ConcurrentMap>> = ConcurrentMap() + + /** + * JSON-response mode only: request ids retired by an inbound cancellation while their POST is + * still awaiting. Treated as settled-without-response so the POST's "a response for every id" + * completion condition is not blocked by a request that will never be answered. + */ + private val cancelledRequestIds: ConcurrentMap = ConcurrentMap() + private val sessionMutex = Mutex() private val streamMutex = Mutex() @@ -241,54 +258,117 @@ public class StreamableHttpServerTransport(private val configuration: Configurat return } + val isTerminated = message is JSONRPCResponse || message is JSONRPCError + + // In JSON mode request-related notifications have no per-request SSE stream. Route them to the + // standalone GET stream before the lookup below, so a notification for an already-retired + // request is delivered there instead of becoming a spurious error. + if (configuration.enableJsonResponse && !isTerminated) { + val standaloneStream = streamsMapping[STANDALONE_SSE_STREAM_ID] + standaloneStream?.let { emitOnStream(STANDALONE_SSE_STREAM_ID, it.session, message) } + return + } + val streamId = requestToStreamMapping[routingRequestId] - ?: error("No connection established for request id $routingRequestId") + if (streamId == null) { + // Request already retired (completed, cancelled, or the JSON POST disconnected). In JSON + // mode a late terminal response is undeliverable and dropped. In SSE mode an unknown id + // is a real routing error. + if (configuration.enableJsonResponse) return + error("No connection established for request id $routingRequestId") + } val activeStream = streamsMapping[streamId] if (!configuration.enableJsonResponse) { activeStream?.let { stream -> emitOnStream(streamId, stream.session, message) } - } - - val isTerminated = message is JSONRPCResponse || message is JSONRPCError - if (!isTerminated) { - if (configuration.enableJsonResponse) { - // In JSON response mode there is no per-request SSE stream, so route notifications - // that are logically associated with a request to the standalone GET SSE stream. - val standaloneStream = streamsMapping[STANDALONE_SSE_STREAM_ID] - standaloneStream?.let { emitOnStream(STANDALONE_SSE_STREAM_ID, it.session, message) } - } - return + if (!isTerminated) return } requestToResponseMapping[responseRequestId!!] = message - val relatedIds = requestToStreamMapping.filterValues { it == streamId }.keys - - if (relatedIds.any { it !in requestToResponseMapping }) return streamMutex.withLock { - if (activeStream == null) error("No connection established for request ID: $routingRequestId") + // The request may have been retired (disconnect or cancellation cleanup) between recording + // the response above and acquiring the lock. Reclaim the stray mapping and drop it. + if (routingRequestId !in requestToStreamMapping) { + requestToResponseMapping.remove(routingRequestId) + return + } + // Recomputed under the lock. Handlers may respond concurrently, and the sender that + // observes the last settled request also retires the ids, so a losing sender sees an + // empty set. A cancelled request counts as settled even though it produces no response. + val relatedIds = requestToStreamMapping.filterValues { it == streamId }.keys + if (relatedIds.isEmpty() || + relatedIds.any { it !in requestToResponseMapping && it !in cancelledRequestIds } + ) { + return + } if (configuration.enableJsonResponse) { - activeStream.call.response.header(HttpHeaders.ContentType, ContentType.Application.Json.toString()) - sessionId?.let { activeStream.call.response.header(MCP_SESSION_ID_HEADER, it) } - val responses = relatedIds.mapNotNull { requestToResponseMapping[it] } - val payload = if (responses.size == 1) { - responses.first() - } else { - responses + // Hand the collected responses to the POST suspended in [handlePostRequest], which + // completes the HTTP exchange. A null pendingPost means the POST was already retired + // (disconnect or cancellation), so the responses are undeliverable and dropped quietly. + val pendingPost = pendingJsonResponses.remove(streamId) + if (pendingPost == null) { + cleanUpRequests(relatedIds) + return } - activeStream.call.respond(payload) + val responses = relatedIds.mapNotNull { requestToResponseMapping[it] } + cleanUpRequests(relatedIds) + pendingPost.complete(responses) } else { - activeStream.session?.close() + // A null stream means the per-request SSE connection closed before the response + // settled, so the ids are retired rather than leaked. + activeStream?.session?.close() + cleanUpRequests(relatedIds) } + } + } - // Clean up - relatedIds.forEach { requestId -> - requestToResponseMapping.remove(requestId) - requestToStreamMapping.remove(requestId) + private fun cleanUpRequests(ids: Collection) { + ids.forEach { requestId -> + requestToResponseMapping.remove(requestId) + requestToStreamMapping.remove(requestId) + cancelledRequestIds.remove(requestId) + } + } + + /** + * Scans a delivered POST body for `notifications/cancelled` and retires each referenced request + * that is awaiting on a JSON-mode POST. No-op unless [Configuration.enableJsonResponse]. + */ + private suspend fun retireCancelledJsonRequests(messages: List) { + if (!configuration.enableJsonResponse) return + val cancelledIds = messages.asSequence() + .filterIsInstance() + .filter { it.method == Method.Defined.NotificationsCancelled.value } + .mapNotNull { notification -> + val params = notification.params ?: return@mapNotNull null + runCatching { + McpJson.decodeFromJsonElement(params).requestId + }.getOrNull() } + .toList() + cancelledIds.forEach { retireCancelledJsonRequest(it) } + } + + /** + * Marks a request cancelled by the peer as settled without a response and completes its POST once + * every request that POST delivered is settled (answered or cancelled). Without this, a cancelled + * request leaves the POST awaiting a response that, by design, is never sent. + */ + private suspend fun retireCancelledJsonRequest(requestId: RequestId) { + streamMutex.withLock { + val streamId = requestToStreamMapping[requestId] ?: return + cancelledRequestIds[requestId] = Unit + val relatedIds = requestToStreamMapping.filterValues { it == streamId }.keys + val allSettled = relatedIds.all { it in requestToResponseMapping || it in cancelledRequestIds } + if (!allSettled) return + val pendingPost = pendingJsonResponses.remove(streamId) ?: return + val responses = relatedIds.mapNotNull { requestToResponseMapping[it] } + cleanUpRequests(relatedIds) + pendingPost.complete(responses) } } @@ -304,6 +384,10 @@ public class StreamableHttpServerTransport(private val configuration: Configurat streamsMapping.clear() requestToStreamMapping.clear() requestToResponseMapping.clear() + cancelledRequestIds.clear() + // Release any POST calls suspended awaiting JSON responses. + pendingJsonResponses.values.forEach { it.cancel() } + pendingJsonResponses.clear() invokeOnCloseCallback() } } @@ -339,7 +423,12 @@ public class StreamableHttpServerTransport(private val configuration: Configurat } /** - * Handles POST requests containing JSON-RPC messages + * Handles POST requests containing JSON-RPC messages. + * + * In JSON-response mode this suspends until a response has been produced for every request in + * the POST body, then completes the HTTP exchange with the collected JSON (a single object for + * one request, an array for a batch). Bodies without requests are acknowledged with 202 + * immediately. In SSE mode the responses are emitted onto the per-request SSE stream instead. */ public suspend fun handlePostRequest(session: ServerSSESession?, call: ApplicationCall) { try { @@ -407,6 +496,9 @@ public class StreamableHttpServerTransport(private val configuration: Configurat if (!hasRequest) { call.respondNullable(status = HttpStatusCode.Accepted, message = null) messages.forEach { message -> _onMessage(message) } + // A cancellation may target a request still awaiting delivery on a JSON-mode POST; + // retire it so that POST completes instead of waiting for a response that never comes. + retireCancelledJsonRequests(messages) return } @@ -429,13 +521,50 @@ public class StreamableHttpServerTransport(private val configuration: Configurat maybeSendPrimingEvent(streamId, session, clientProtocolVersion) } + val pendingPost = if (configuration.enableJsonResponse) { + CompletableDeferred>() + } else { + null + } + streamMutex.withLock { streamsMapping[streamId] = SessionContext(session, call) messages.filterIsInstance().forEach { requestToStreamMapping[it.id] = streamId } + pendingPost?.let { pendingJsonResponses[streamId] = it } } call.coroutineContext.job.invokeOnCompletion { streamsMapping.remove(streamId) } - messages.forEach { message -> _onMessage(message) } + if (pendingPost == null) { + // SSE mode: responses are emitted onto this call's SSE stream by send(). + messages.forEach { message -> _onMessage(message) } + return + } + + // JSON-response mode. Handlers may run concurrently, so returning from delivery no longer + // implies the responses were produced. Suspend until send() has collected a response for + // every request in this POST. A client disconnect or transport close cancels this call. + try { + messages.forEach { message -> _onMessage(message) } + retireCancelledJsonRequests(messages) + val responses = pendingPost.await() + sessionId?.let { call.response.header(MCP_SESSION_ID_HEADER, it) } + if (responses.isEmpty()) { + // Every request in this POST was cancelled before producing a response; there is + // nothing to return, so acknowledge with 202 rather than an empty JSON body. + call.respondNullable(status = HttpStatusCode.Accepted, message = null) + } else { + call.response.header(HttpHeaders.ContentType, ContentType.Application.Json.toString()) + val payload = if (responses.size == 1) responses.first() else responses + call.respond(payload) + } + } finally { + // On cancellation or failure, retire this POST's ids so a late response is dropped + // quietly at the send() boundary instead of leaking mappings. Idempotent no-ops on + // the success path (send() or retirement already cleaned up). + pendingJsonResponses.remove(streamId) + pendingPost.cancel() + cleanUpRequests(messages.filterIsInstance().map { it.id }) + } } catch (e: CancellationException) { throw e } catch (e: Exception) { diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransportTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransportTest.kt index d7dec4e..172889a 100644 --- a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransportTest.kt +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransportTest.kt @@ -19,14 +19,19 @@ import io.ktor.http.HttpHeaders import io.ktor.http.HttpStatusCode import io.ktor.http.contentType import io.ktor.serialization.kotlinx.json.json +import io.ktor.server.application.ApplicationCall import io.ktor.server.application.install import io.ktor.server.routing.post import io.ktor.server.routing.routing +import io.ktor.server.sse.ServerSSESession import io.ktor.server.testing.ApplicationTestBuilder import io.ktor.server.testing.testApplication +import io.ktor.sse.ServerSentEvent import io.ktor.utils.io.ByteReadChannel import io.ktor.utils.io.readLine import io.ktor.utils.io.readUTF8Line +import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotification +import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult import io.modelcontextprotocol.kotlin.sdk.types.Implementation @@ -46,6 +51,16 @@ import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities import io.modelcontextprotocol.kotlin.sdk.types.Tool import io.modelcontextprotocol.kotlin.sdk.types.ToolSchema import io.modelcontextprotocol.kotlin.sdk.types.toJSON +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.async +import kotlinx.coroutines.cancel +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withTimeout import kotlinx.serialization.builtins.ListSerializer import kotlinx.serialization.json.buildJsonObject import kotlinx.serialization.json.put @@ -53,11 +68,13 @@ import org.junit.jupiter.params.ParameterizedTest import org.junit.jupiter.params.provider.Arguments import org.junit.jupiter.params.provider.MethodSource import java.util.concurrent.atomic.AtomicBoolean +import kotlin.coroutines.CoroutineContext import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertFalse import kotlin.test.assertNotNull +import kotlin.time.Duration.Companion.seconds import io.ktor.client.plugins.contentnegotiation.ContentNegotiation as ClientContentNegotiation import io.ktor.server.plugins.contentnegotiation.ContentNegotiation as ServerContentNegotiation @@ -349,6 +366,306 @@ class StreamableHttpServerTransportTest { assertEquals("second", secondMeta?.get("label")?.jsonPrimitive?.content)*/ } + @Test + fun `json response waits for a response produced after message delivery returns`() = testApplication { + configTestServer() + + val client = createTestClient() + + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = true), + ) + // Simulates concurrent dispatch: message delivery returns before the handler responds. + val handlerScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + try { + transport.onMessage { message -> + if (message is JSONRPCRequest) { + if (message.method == Method.Defined.Initialize.value) { + transport.send(JSONRPCResponse(message.id, EmptyResult())) + } else { + handlerScope.launch { + delay(100) + transport.send(JSONRPCResponse(message.id, EmptyResult())) + } + } + } + } + + configureTransportEndpoint(transport) + + val initResponse = client.post(path) { + addStreamableHeaders() + setBody(buildInitializeRequestPayload()) + } + initResponse.status shouldBe HttpStatusCode.OK + + val request = JSONRPCRequest(id = RequestId("async-1"), method = Method.Defined.ToolsList.value) + val response = withTimeout(10.seconds) { + client.post(path) { + addStreamableHeaders() + header(MCP_SESSION_ID_HEADER, initResponse.headers[MCP_SESSION_ID_HEADER]) + setBody(encodeMessages(listOf(request))) + } + } + + response.status shouldBe HttpStatusCode.OK + response.body() shouldBe JSONRPCResponse(request.id) + } finally { + handlerScope.cancel() + } + } + + @Test + fun `json response for a batch waits until every asynchronously produced response is ready`() = testApplication { + configTestServer() + + val client = createTestClient() + + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = true), + ) + val firstRequest = JSONRPCRequest(id = RequestId("slow"), method = Method.Defined.ToolsList.value) + val secondRequest = JSONRPCRequest(id = RequestId("fast"), method = Method.Defined.ResourcesList.value) + + // Simulates concurrent dispatch with out-of-order completion: the second request responds + // first; the POST must stay open until the slower first response settles too. + val handlerScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + try { + transport.onMessage { message -> + if (message is JSONRPCRequest) { + if (message.method == Method.Defined.Initialize.value) { + transport.send(JSONRPCResponse(message.id, EmptyResult())) + } else { + handlerScope.launch { + delay(if (message.id == firstRequest.id) 200 else 50) + transport.send(JSONRPCResponse(message.id, EmptyResult())) + } + } + } + } + + configureTransportEndpoint(transport) + + val initResponse = client.post(path) { + addStreamableHeaders() + setBody(buildInitializeRequestPayload()) + } + initResponse.status shouldBe HttpStatusCode.OK + + val response = withTimeout(10.seconds) { + client.post(path) { + addStreamableHeaders() + header(MCP_SESSION_ID_HEADER, initResponse.headers[MCP_SESSION_ID_HEADER]) + setBody(encodeMessages(listOf(firstRequest, secondRequest))) + } + } + + response.status shouldBe HttpStatusCode.OK + val responses = response.body>() + responses.map { it.id }.toSet() shouldBe setOf(firstRequest.id, secondRequest.id) + } finally { + handlerScope.cancel() + } + } + + @Test + fun `json response POST completes with 202 when its only request is cancelled`() = testApplication { + configTestServer() + + val client = createTestClient() + + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = true), + ) + val requestDelivered = CompletableDeferred() + // Concurrent dispatch: delivery returns before a response is produced. This request is never + // answered here, modelling a handler cancelled by notifications/cancelled. + transport.onMessage { message -> + if (message is JSONRPCRequest) { + if (message.method == Method.Defined.Initialize.value) { + transport.send(JSONRPCResponse(message.id, EmptyResult())) + } else { + requestDelivered.complete(Unit) + } + } + } + + configureTransportEndpoint(transport) + + val clientScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + try { + val initResponse = client.post(path) { + addStreamableHeaders() + setBody(buildInitializeRequestPayload()) + } + initResponse.status shouldBe HttpStatusCode.OK + val sessionId = initResponse.headers[MCP_SESSION_ID_HEADER] + + val request = JSONRPCRequest(id = RequestId("cancel-me"), method = Method.Defined.ToolsList.value) + val requestPost = clientScope.async { + client.post(path) { + addStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody(encodeMessages(listOf(request))) + } + } + + // The request must be registered in-flight before it can be cancelled. + withTimeout(5.seconds) { requestDelivered.await() } + + val cancellation = CancelledNotification( + CancelledNotificationParams(requestId = request.id, reason = "client cancelled"), + ) + val cancelResponse = client.post(path) { + addStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody(encodeMessages(listOf(cancellation.toJSON()))) + } + cancelResponse.status shouldBe HttpStatusCode.Accepted + + // Without retirement of the cancelled id this await never completes and the timeout fires. + val response = withTimeout(10.seconds) { requestPost.await() } + response.status shouldBe HttpStatusCode.Accepted + } finally { + clientScope.cancel() + } + } + + @Test + fun `json response for a batch returns only the non-cancelled responses`() = testApplication { + configTestServer() + + val client = createTestClient() + + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = true), + ) + val answeredA = JSONRPCRequest(id = RequestId("a"), method = Method.Defined.ToolsList.value) + val answeredB = JSONRPCRequest(id = RequestId("b"), method = Method.Defined.ToolsList.value) + val cancelledC = JSONRPCRequest(id = RequestId("c"), method = Method.Defined.ToolsList.value) + + val cancelledDelivered = CompletableDeferred() + val handlerScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + try { + transport.onMessage { message -> + if (message is JSONRPCRequest) { + when { + message.method == Method.Defined.Initialize.value -> + transport.send(JSONRPCResponse(message.id, EmptyResult())) + + message.id == answeredA.id || message.id == answeredB.id -> + handlerScope.launch { transport.send(JSONRPCResponse(message.id, EmptyResult())) } + + message.id == cancelledC.id -> cancelledDelivered.complete(Unit) + } + } + } + + configureTransportEndpoint(transport) + + val initResponse = client.post(path) { + addStreamableHeaders() + setBody(buildInitializeRequestPayload()) + } + initResponse.status shouldBe HttpStatusCode.OK + val sessionId = initResponse.headers[MCP_SESSION_ID_HEADER] + + val batchPost = handlerScope.async { + client.post(path) { + addStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody(encodeMessages(listOf(answeredA, answeredB, cancelledC))) + } + } + + withTimeout(5.seconds) { cancelledDelivered.await() } + + val cancellation = CancelledNotification(CancelledNotificationParams(requestId = cancelledC.id)) + client.post(path) { + addStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody(encodeMessages(listOf(cancellation.toJSON()))) + }.status shouldBe HttpStatusCode.Accepted + + val response = withTimeout(10.seconds) { batchPost.await() } + response.status shouldBe HttpStatusCode.OK + val responses = response.body>() + responses.map { it.id }.toSet() shouldBe setOf(answeredA.id, answeredB.id) + } finally { + handlerScope.cancel() + } + } + + @Test + fun `json response send for a retired request id is dropped without error`() = runTest { + val transport = StreamableHttpServerTransport(enableJsonResponse = true) + // No stream is mapped for this id (already retired by disconnect/cancellation cleanup); + // a late terminal response must be dropped quietly, not raise "No connection established". + transport.send(JSONRPCResponse(RequestId("retired"), EmptyResult())) + } + + @Test + fun `sse response send for an unknown request id still errors`() = runTest { + val transport = StreamableHttpServerTransport(enableJsonResponse = false) + assertFailsWith { + transport.send(JSONRPCResponse(RequestId("unknown"), EmptyResult())) + } + } + + @Test + fun `sse terminal response is dropped without leaking when the per-request stream is gone`() = testApplication { + configTestServer() + val client = createTestClient() + + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = false), + ) + transport.setSessionIdGenerator(null) // stateless: accept requests without an init handshake + val requestId = RequestId("held-request") + val delivered = CompletableDeferred() + val release = CompletableDeferred() + transport.onMessage { message -> + if (message is JSONRPCRequest) { + delivered.complete(Unit) + release.await() // keep the POST and its stream mapping alive until the test releases it + } + } + + application { + routing { + post(path) { + transport.handlePostRequest(FakeServerSSESession(call, call.coroutineContext), call) + } + } + } + + val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + try { + val post = scope.async { + client.post(path) { + addStreamableHeaders() + setBody( + encodeMessages( + listOf(JSONRPCRequest(id = requestId, method = Method.Defined.ToolsList.value)), + ), + ) + } + } + withTimeout(5.seconds) { delivered.await() } + + // Drop the per-request SSE stream while the request is still mapped in-flight; the handler's + // terminal response is now undeliverable and must be dropped quietly rather than raise + // "No connection established" (which would also strand the request mappings). + transport.closeSseStream(requestId) + transport.send(JSONRPCResponse(requestId, EmptyResult())) + + post.cancel() + } finally { + release.complete(Unit) + scope.cancel() + } + } + @ParameterizedTest @MethodSource("invalidPayloads") fun `POST with a null or empty body returns an error`(payload: String?) = testApplication { @@ -694,3 +1011,12 @@ class StreamableHttpServerTransportTest { } } } + +/** No-op [ServerSSESession] double that lets a POST route drive [StreamableHttpServerTransport] in SSE mode. */ +private class FakeServerSSESession( + override val call: ApplicationCall, + override val coroutineContext: CoroutineContext, +) : ServerSSESession { + override suspend fun send(event: ServerSentEvent) {} + override suspend fun close() {} +}