diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/WebSocketMessageSizeIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/WebSocketMessageSizeIntegrationTest.kt new file mode 100644 index 000000000..564bc7654 --- /dev/null +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/WebSocketMessageSizeIntegrationTest.kt @@ -0,0 +1,69 @@ +package io.modelcontextprotocol.kotlin.sdk.integration + +import io.kotest.matchers.shouldBe +import io.ktor.client.HttpClient +import io.ktor.client.plugins.websocket.WebSockets +import io.ktor.client.plugins.websocket.webSocketSession +import io.ktor.client.request.header +import io.ktor.http.HttpHeaders +import io.ktor.server.engine.embeddedServer +import io.ktor.websocket.Frame +import io.modelcontextprotocol.kotlin.sdk.client.Client +import io.modelcontextprotocol.kotlin.sdk.client.mcpWebSocketTransport +import io.modelcontextprotocol.kotlin.sdk.server.Server +import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions +import io.modelcontextprotocol.kotlin.sdk.server.mcpWebSocket +import io.modelcontextprotocol.kotlin.sdk.shared.MCP_SUBPROTOCOL +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.test.utils.actualPort +import io.modelcontextprotocol.kotlin.test.utils.runIntegrationTest +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.put +import kotlinx.serialization.json.putJsonObject +import org.junit.jupiter.api.Test +import kotlin.time.Duration.Companion.seconds +import io.ktor.client.engine.cio.CIO as ClientCIO +import io.ktor.server.cio.CIO as ServerCIO + +class WebSocketMessageSizeIntegrationTest { + + @Test + fun `server closes a connection that sends a message over 16 MiB and keeps serving others`() = + runIntegrationTest(timeout = 30.seconds) { + val server = Server(Implementation("test-server", "1.0.0"), ServerOptions(ServerCapabilities())) + val engine = embeddedServer(ServerCIO, host = HOST, port = 0) { mcpWebSocket { server } } + .startSuspend(wait = false) + val httpClient = HttpClient(ClientCIO) { install(WebSockets) } + val url = "ws://$HOST:${engine.actualPort()}" + try { + val peer = httpClient.webSocketSession(url) { + header(HttpHeaders.SecWebSocketProtocol, MCP_SUBPROTOCOL) + } + + // A valid request: without the limit the server would answer it instead of closing. + peer.send(Frame.Text(ping(padding = "a".repeat(16 * 1024 * 1024)))) + + peer.incoming.receiveCatching().getOrNull() shouldBe null + val client = Client(Implementation("test-client", "1.0.0")) + client.connect(httpClient.mcpWebSocketTransport(url)) + client.ping() + client.close() + } finally { + httpClient.close() + engine.stopSuspend(gracePeriodMillis = 0, timeoutMillis = 1_000) + server.close() + } + } + + private fun ping(padding: String): String = buildJsonObject { + put("jsonrpc", "2.0") + put("id", 1) + put("method", "ping") + putJsonObject("params") { put("padding", padding) } + }.toString() + + private companion object { + const val HOST = "127.0.0.1" + } +} diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt index 820864800..2d8ace42f 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransport.kt @@ -64,12 +64,19 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { } catch (_: ClosedReceiveChannelException) { logger.debug { "Closed receive channel, exiting" } return@launch + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + // The session is already failing (e.g. FrameTooBigException) and Ktor closes it. + invokeOnErrorCallback(e) + return@launch } + // Keep reading after a bad message: an open session nobody reads buffers frames without limit. if (message !is Frame.Text) { val e = IllegalArgumentException("Expected text frame, got ${message::class.simpleName}: $message") - _onError.invoke(e) - throw e + invokeOnErrorCallback(e) + continue } try { @@ -78,8 +85,7 @@ public abstract class WebSocketMcpTransport : AbstractTransport() { } catch (e: CancellationException) { throw e } catch (e: Exception) { - _onError.invoke(e) - throw e + invokeOnErrorCallback(e) } } } diff --git a/kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransportTest.kt b/kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransportTest.kt new file mode 100644 index 000000000..5c8d25aa7 --- /dev/null +++ b/kotlin-sdk-core/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/WebSocketMcpTransportTest.kt @@ -0,0 +1,75 @@ +package io.modelcontextprotocol.kotlin.sdk.shared + +import io.kotest.matchers.collections.shouldHaveSize +import io.kotest.matchers.shouldBe +import io.kotest.matchers.types.shouldBeInstanceOf +import io.ktor.websocket.Frame +import io.ktor.websocket.FrameTooBigException +import io.ktor.websocket.WebSocketExtension +import io.ktor.websocket.WebSocketSession +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Job +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.channels.ReceiveChannel +import kotlinx.coroutines.channels.SendChannel +import kotlinx.coroutines.test.runTest +import org.junit.jupiter.api.Test +import java.util.concurrent.CopyOnWriteArrayList +import kotlin.coroutines.CoroutineContext + +class WebSocketMcpTransportTest { + + @Test + fun `should keep reading after messages it cannot parse`() = runTest { + val session = FakeWebSocketSession() + val transport = TestTransport(session) + val errors = CopyOnWriteArrayList() + val received = CompletableDeferred() + transport.onError { errors += it } + transport.onMessage { received.complete(it) } + transport.start() + + session.incomingFrames.send(Frame.Binary(true, byteArrayOf(1, 2, 3))) + session.incomingFrames.send(Frame.Text("{")) + session.incomingFrames.send(Frame.Text("""{"jsonrpc":"2.0","id":1,"method":"ping"}""")) + + received.await().shouldBeInstanceOf().method shouldBe "ping" + errors shouldHaveSize 2 + session.incomingFrames.close() + } + + @Test + fun `should report a session failure through onError`() = runTest { + val session = FakeWebSocketSession() + val transport = TestTransport(session) + val error = CompletableDeferred() + transport.onError { error.complete(it) } + transport.start() + + session.incomingFrames.close(FrameTooBigException(42)) + + error.await().shouldBeInstanceOf() + } + + private class TestTransport(override val session: WebSocketSession) : WebSocketMcpTransport() { + override suspend fun initializeSession() = Unit + } +} + +private class FakeWebSocketSession : WebSocketSession { + val incomingFrames = Channel(Channel.UNLIMITED) + + override val coroutineContext: CoroutineContext = Job() + override var masking: Boolean = false + override var maxFrameSize: Long = Long.MAX_VALUE + override val incoming: ReceiveChannel get() = incomingFrames + override val outgoing: SendChannel = Channel(Channel.UNLIMITED) + override val extensions: List> = emptyList() + + override suspend fun flush() = Unit + + @Deprecated("Use cancel() instead.", level = DeprecationLevel.ERROR) + override fun terminate() = Unit +} diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/WebSocketMcpKtorServerExtensions.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/WebSocketMcpKtorServerExtensions.kt index a05daefa5..e4957f57a 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/WebSocketMcpKtorServerExtensions.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/WebSocketMcpKtorServerExtensions.kt @@ -18,8 +18,13 @@ import kotlinx.coroutines.awaitCancellation private val logger = KotlinLogging.logger {} +// Ktor's default maxFrameSize does not limit inbound messages at all. +private const val DEFAULT_MAX_FRAME_SIZE: Long = 16L * 1024 * 1024 + /** * Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket. + * + * Set `maxFrameSize` when installing the `WebSockets` plugin: by default Ktor does not limit inbound messages. */ @KtorDsl public fun Route.mcpWebSocket(block: () -> Server) { @@ -30,6 +35,8 @@ public fun Route.mcpWebSocket(block: () -> Server) { /** * Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket. + * + * Set `maxFrameSize` when installing the `WebSockets` plugin: by default Ktor does not limit inbound messages. */ @KtorDsl public fun Route.mcpWebSocket(path: String, block: () -> Server) { @@ -40,10 +47,14 @@ public fun Route.mcpWebSocket(path: String, block: () -> Server) { /** * Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket. + * + * Installs the `WebSockets` plugin with inbound messages limited to 16 MiB. */ @KtorDsl public fun Application.mcpWebSocket(block: () -> Server) { - install(WebSockets) + install(WebSockets) { + maxFrameSize = DEFAULT_MAX_FRAME_SIZE + } routing { mcpWebSocket(block) @@ -52,10 +63,14 @@ public fun Application.mcpWebSocket(block: () -> Server) { /** * Configures the Ktor Application to handle Model Context Protocol (MCP) over WebSocket at the specified path. + * + * Installs the `WebSockets` plugin with inbound messages limited to 16 MiB. */ @KtorDsl public fun Application.mcpWebSocket(path: String, block: () -> Server) { - install(WebSockets) + install(WebSockets) { + maxFrameSize = DEFAULT_MAX_FRAME_SIZE + } routing { mcpWebSocket(path, block)