diff --git a/README.md b/README.md index aefadcf6a..a08b8016d 100644 --- a/README.md +++ b/README.md @@ -803,6 +803,48 @@ embeddedServer(CIO, port = 3000) { --> +**Session limits:** `mcpStreamableHttp()` keeps every session in memory until the client deletes it, the server closes +it, or it expires. A session expires after `sessionIdleTimeout` (30 minutes by default) without requests; its client +then receives `404 Not Found` and has to initialize a new session. An open GET stream does not keep a session alive +(expiry closes the stream too), so a client that only listens for notifications should send `ping` requests. At most +`maxSessions` sessions (10,000 by default) are open at once; while that many are open, a request that would open another +one receives `503 Service Unavailable`. Any request other than `initialize` that carries no `Mcp-Session-Id`, a plain +`GET /mcp` included, receives `400 Bad Request`, so point health checks at a separate route. Size both limits to your +traffic: + + + + +```kotlin +embeddedServer(CIO, port = 3000) { + mcpStreamableHttp(sessionIdleTimeout = 10.minutes, maxSessions = 1_000) { + MyServer() + } +}.start(wait = true) +``` + + + + **CORS for browser-based clients (e.g. MCP Inspector):** if you connect from a browser-based client you need to install the Ktor CORS plugin so that MCP-specific headers are allowed and exposed: @@ -869,7 +911,7 @@ embeddedServer(CIO, port = 3000) { - + Prefer Streamable HTTP for new projects. diff --git a/conformance-test/run-conformance.sh b/conformance-test/run-conformance.sh index 279254ebb..0c22b2f7b 100755 --- a/conformance-test/run-conformance.sh +++ b/conformance-test/run-conformance.sh @@ -44,7 +44,7 @@ start_server() { echo "Waiting for server to be ready..." local retries=0 local max_retries=30 - while ! curl -sf "$SERVER_URL" > /dev/null 2>&1; do + while ! curl -s -o /dev/null "$SERVER_URL" 2>&1; do retries=$((retries + 1)) if [ "$retries" -ge "$max_retries" ]; then echo "ERROR: Server failed to start after $max_retries attempts" diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ClientDisconnectIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ClientDisconnectIntegrationTest.kt new file mode 100644 index 000000000..e84a10ce0 --- /dev/null +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/ClientDisconnectIntegrationTest.kt @@ -0,0 +1,239 @@ +package io.modelcontextprotocol.kotlin.sdk.integration + +import io.kotest.assertions.nondeterministic.eventually +import io.kotest.matchers.maps.shouldBeEmpty +import io.kotest.matchers.maps.shouldHaveSize +import io.kotest.matchers.nulls.shouldNotBeNull +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldStartWith +import io.ktor.client.HttpClient +import io.ktor.client.request.header +import io.ktor.client.request.post +import io.ktor.client.request.setBody +import io.ktor.http.ContentType +import io.ktor.http.HttpHeaders +import io.ktor.http.HttpMethod +import io.ktor.http.HttpStatusCode +import io.ktor.http.contentType +import io.ktor.server.application.Application +import io.ktor.server.application.createApplicationPlugin +import io.ktor.server.application.install +import io.ktor.server.engine.embeddedServer +import io.ktor.server.request.httpMethod +import io.ktor.server.routing.get +import io.ktor.server.routing.routing +import io.modelcontextprotocol.kotlin.sdk.server.Server +import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions +import io.modelcontextprotocol.kotlin.sdk.server.mcp +import io.modelcontextprotocol.kotlin.sdk.server.mcpStreamableHttp +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.delay +import kotlinx.coroutines.job +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeout +import java.net.Socket +import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds +import io.ktor.client.engine.cio.CIO as ClientCIO +import io.ktor.server.cio.CIO as ServerCIO +import io.ktor.server.sse.SSE as ServerSSE + +/** + * Runs against a real CIO server, because what matters here is engine behavior: CIO never cancels a handler that is + * only waiting after its client disconnects, so a client that goes away without a word must not leave its session + * behind. The clients are raw sockets, so they can disconnect without saying anything. + */ +class ClientDisconnectIntegrationTest { + + private fun testServer() = Server(Implementation("test-server", "1.0"), ServerOptions(ServerCapabilities())) + + @Test + fun `an SSE session ends when its client disconnects`(): Unit = runBlocking { + val server = testServer() + withCioServer({ + install(ServerSSE) + routing { mcp { server } } + }) { port -> + val sessionId = connect(port, LEGACY_SSE_REQUEST).use { socket -> + socket.readLineMatching { "sessionId=" in it }.substringAfter("sessionId=") + } + + eventually(10.seconds) { server.sessions.shouldBeEmpty() } + HttpClient(ClientCIO).use { client -> + client.post("http://$HOST:$port/?sessionId=$sessionId") { + contentType(ContentType.Application.Json) + setBody("""{"jsonrpc":"2.0","id":1,"method":"ping"}""") + }.status shouldBe HttpStatusCode.NotFound + } + } + } + + @Test + fun `an SSE handler ends when the server closes its session`(): Unit = runBlocking { + val server = testServer() + val finishedGets = AtomicInteger() + withCioServer({ + countFinishedGets(finishedGets) + install(ServerSSE) + routing { mcp { server } } + }) { port -> + connect(port, LEGACY_SSE_REQUEST).use { socket -> + socket.readLineMatching { "sessionId=" in it } + eventually(10.seconds) { server.sessions shouldHaveSize 1 } + + server.sessions.values.single().close() + + // The client stays connected, so only the handler returning can finish the call. + eventually(10.seconds) { finishedGets.get() shouldBe 1 } + } + } + } + + @Test + fun `an application route next to MCP runs to completion after its client disconnects`(): Unit = runBlocking { + val handlerStarted = CompletableDeferred() + val handlerCompleted = CompletableDeferred() + withCioServer({ + mcp { testServer() } + routing { + get("/slow") { + handlerStarted.complete(Unit) + try { + delay(2.seconds) + handlerCompleted.complete(true) + } catch (e: CancellationException) { + handlerCompleted.complete(false) + throw e + } + } + } + }) { port -> + connect(port, "GET /slow HTTP/1.1\r\nHost: localhost\r\n\r\n").use { + withTimeout(10.seconds) { handlerStarted.await() } + } + + withTimeout(10.seconds) { handlerCompleted.await() } shouldBe true + } + } + + @Test + fun `a GET stream ends as soon as its client disconnects`(): Unit = runBlocking { + val server = testServer() + val finishedGets = AtomicInteger() + withCioServer({ + countFinishedGets(finishedGets) + mcpStreamableHttp { server } + }) { port -> + val sessionId = initializeSession(port) + val stream = openStream(port, sessionId) + withContext(Dispatchers.IO) { stream.close() } + + // Long before the session could expire. + eventually(10.seconds) { finishedGets.get() shouldBe 1 } + server.sessions shouldHaveSize 1 + } + } + + @Test + fun `a POST whose client disconnects does not keep its session from expiring`(): Unit = runBlocking { + val toolStarted = CompletableDeferred() + val server = Server( + Implementation("test-server", "1.0"), + ServerOptions(ServerCapabilities(tools = ServerCapabilities.Tools())), + ) { + addTool(name = "wait", description = "Never answers") { + toolStarted.complete(Unit) + awaitCancellation() + } + } + withCioServer({ mcpStreamableHttp(sessionIdleTimeout = 1.seconds) { server } }) { port -> + val sessionId = initializeSession(port) + val body = """{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"wait","arguments":{}}}""" + val request = "POST /mcp HTTP/1.1\r\nHost: localhost\r\nAccept: application/json, text/event-stream\r\n" + + "Content-Type: application/json\r\nMcp-Session-Id: $sessionId\r\n" + + "Content-Length: ${body.toByteArray().size}\r\n\r\n$body" + connect(port, request).use { + withTimeout(10.seconds) { toolStarted.await() } + } + + // The request is still unanswered, but with its client gone it no longer holds the session. + eventually(10.seconds) { server.sessions.shouldBeEmpty() } + } + } + + private suspend fun initializeSession(port: Int): String = HttpClient(ClientCIO).use { client -> + val sessionId = client.post("http://$HOST:$port/mcp") { + header(HttpHeaders.Accept, "${ContentType.Application.Json}, ${ContentType.Text.EventStream}") + contentType(ContentType.Application.Json) + setBody(INITIALIZE_REQUEST) + }.headers["mcp-session-id"].shouldNotBeNull() + client.post("http://$HOST:$port/mcp") { + header(HttpHeaders.Accept, "${ContentType.Application.Json}, ${ContentType.Text.EventStream}") + header("mcp-session-id", sessionId) + contentType(ContentType.Application.Json) + setBody("""{"jsonrpc":"2.0","method":"notifications/initialized"}""") + }.status shouldBe HttpStatusCode.Accepted + sessionId + } + + /** Opens the session's GET stream on a raw socket, which the caller can drop without a word. */ + private suspend fun openStream(port: Int, sessionId: String): Socket { + val socket = connect( + port, + "GET /mcp HTTP/1.1\r\nHost: localhost\r\nAccept: text/event-stream\r\nMcp-Session-Id: $sessionId\r\n\r\n", + ) + socket.readLineMatching { it.startsWith("HTTP/") } shouldStartWith "HTTP/1.1 200" + return socket + } + + /** Opens a raw connection and sends [request] on it, so that the test can later drop it without a word. */ + private suspend fun connect(port: Int, request: String): Socket = withContext(Dispatchers.IO) { + Socket(HOST, port).apply { getOutputStream().write(request.toByteArray()) } + } + + /** Reads this connection until a line matches [predicate], and returns that line. */ + private suspend fun Socket.readLineMatching(predicate: (String) -> Boolean): String = withContext(Dispatchers.IO) { + val lines = getInputStream().bufferedReader() + generateSequence { lines.readLine() }.first(predicate) + } + + /** + * Counts the GET calls that have ended, whether their handler returned or they were cancelled. A cancelled call + * never finishes sending its response, so its end shows only as the completion of its job. + */ + private fun Application.countFinishedGets(counter: AtomicInteger) { + install( + createApplicationPlugin("FinishedGets") { + onCall { call -> + if (call.request.httpMethod == HttpMethod.Get) { + call.coroutineContext.job.invokeOnCompletion { counter.incrementAndGet() } + } + } + }, + ) + } + + private suspend fun withCioServer(module: Application.() -> Unit, block: suspend (port: Int) -> Unit) { + val server = embeddedServer(ServerCIO, host = HOST, port = 0, module = module).start(wait = false) + try { + block(server.engine.resolvedConnectors().first().port) + } finally { + server.stop(0, 0) + } + } + + private companion object { + const val HOST = "127.0.0.1" + const val LEGACY_SSE_REQUEST = "GET / HTTP/1.1\r\nHost: localhost\r\nAccept: text/event-stream\r\n\r\n" + const val INITIALIZE_REQUEST = + """{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25",""" + + """"capabilities":{},"clientInfo":{"name":"test-client","version":"1.0"}}}""" + } +} diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpClientTransportIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpClientTransportIntegrationTest.kt index 1738ac5a0..91b9e5b82 100644 --- a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpClientTransportIntegrationTest.kt +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpClientTransportIntegrationTest.kt @@ -17,14 +17,20 @@ import io.ktor.server.routing.post import io.ktor.server.routing.routing import io.modelcontextprotocol.kotlin.sdk.client.Client import io.modelcontextprotocol.kotlin.sdk.client.StreamableHttpClientTransport +import io.modelcontextprotocol.kotlin.sdk.client.StreamableHttpError +import io.modelcontextprotocol.kotlin.sdk.server.Server +import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions +import io.modelcontextprotocol.kotlin.sdk.server.mcpStreamableHttp import io.modelcontextprotocol.kotlin.sdk.types.Implementation import io.modelcontextprotocol.kotlin.sdk.types.LoggingMessageNotification import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities import io.modelcontextprotocol.kotlin.test.utils.actualPort import io.modelcontextprotocol.kotlin.test.utils.runIntegrationTest import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.awaitCancellation import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.withTimeout import kotlinx.serialization.json.Json import kotlinx.serialization.json.jsonObject import kotlinx.serialization.json.jsonPrimitive @@ -96,6 +102,30 @@ class StreamableHttpClientTransportIntegrationTest { client.ping() } + @Test + fun `client reports an expired session through onError`() = runIntegrationTest(timeout = 20.seconds) { + val server = Server(Implementation("test-server", "1.0.0"), ServerOptions(ServerCapabilities())) + val ktorServer = embeddedServer(ServerCIO, host = HOST, port = 0) { + mcpStreamableHttp(sessionIdleTimeout = 1.seconds) { server } + }.startSuspend(wait = false) + val httpClient = HttpClient(ClientCIO) { install(SSE) } + val client = Client(Implementation("test-client", "1.0.0")) + try { + val url = "http://$HOST:${ktorServer.actualPort()}$MCP_PATH" + val transport = StreamableHttpClientTransport(httpClient, url) + val sessionNotFound = CompletableDeferred() + transport.onError { if (it is StreamableHttpError && it.code == 404) sessionNotFound.complete(it) } + client.connect(transport) + + // The open GET stream does not keep the session alive, so its reconnect after expiry gets 404. + withTimeout(10.seconds) { sessionNotFound.await() } + } finally { + client.close() + httpClient.close() + ktorServer.stopSuspend(1000, 2000) + } + } + private fun withConnectedClient( getRoute: Route.() -> Unit, configure: Client.() -> Unit = {}, diff --git a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt index 6df352abe..af2813401 100644 --- a/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt +++ b/kotlin-sdk-client/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport.kt @@ -90,6 +90,9 @@ private sealed interface ConnectResult { * Supports automatic SSE reconnection with exponential backoff, stream resumption via the * `Last-Event-ID` header, and explicit session termination. * + * A session the server has terminated answers `404`, reported through [onError] as a [StreamableHttpError]. + * The transport does not start a new session; connect the client again with a new transport. + * * @param client Ktor HTTP client used for all requests * @param url MCP endpoint URL * @param reconnectionOptions reconnection backoff and retry-limit settings for the SSE stream @@ -378,6 +381,13 @@ public class StreamableHttpClientTransport( val responseContentType = e.response?.contentType() return when { + responseStatus == HttpStatusCode.NotFound && sessionId != null -> { + val error = StreamableHttpError(responseStatus.value, "Session $sessionId not found, stream disabled") + logger.warn { error.message } + _onError(error) + true + } + responseStatus == HttpStatusCode.NotFound || responseStatus == HttpStatusCode.MethodNotAllowed -> { logger.info { "Server returned ${responseStatus.value} for GET/SSE, stream disabled." } true diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index 4e5673b20..acfe1ae39 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -61,8 +61,8 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/KtorServerKt { public static final fun mcpStatelessStreamableHttp (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lkotlin/jvm/functions/Function1;)V public static synthetic fun mcpStatelessStreamableHttp$default (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun mcpStatelessStreamableHttp$default (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V - public static final fun mcpStreamableHttp (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)V - public static synthetic fun mcpStreamableHttp$default (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public static final fun mcpStreamableHttp-u4uj9nY (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;Lkotlin/jvm/functions/Function1;JILkotlin/jvm/functions/Function1;)V + public static synthetic fun mcpStreamableHttp-u4uj9nY$default (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;Lkotlin/jvm/functions/Function1;JILkotlin/jvm/functions/Function1;ILjava/lang/Object;)V } public final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt : io/modelcontextprotocol/kotlin/sdk/server/Feature { diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.klib.api b/kotlin-sdk-server/api/kotlin-sdk-server.klib.api index f23c22f66..7d64dd05d 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.klib.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.klib.api @@ -286,7 +286,7 @@ final val io.modelcontextprotocol.kotlin.sdk.server/DnsRebindingProtection // io final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcp(kotlin/Boolean = ..., kotlin.collections/List? = ..., kotlin.collections/List? = ..., kotlin/Long = ..., kotlin/Function1) // io.modelcontextprotocol.kotlin.sdk.server/mcp|mcp@io.ktor.server.application.Application(kotlin.Boolean;kotlin.collections.List?;kotlin.collections.List?;kotlin.Long;kotlin.Function1){}[0] final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcpStatelessStreamableHttp(kotlin/String = ..., kotlin/Boolean = ..., kotlin.collections/List? = ..., kotlin.collections/List? = ..., io.modelcontextprotocol.kotlin.sdk.server/EventStore?, kotlin/Function1) // io.modelcontextprotocol.kotlin.sdk.server/mcpStatelessStreamableHttp|mcpStatelessStreamableHttp@io.ktor.server.application.Application(kotlin.String;kotlin.Boolean;kotlin.collections.List?;kotlin.collections.List?;io.modelcontextprotocol.kotlin.sdk.server.EventStore?;kotlin.Function1){}[0] final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcpStatelessStreamableHttp(kotlin/String = ..., kotlin/Boolean = ..., kotlin.collections/List? = ..., kotlin.collections/List? = ..., kotlin/Function1) // io.modelcontextprotocol.kotlin.sdk.server/mcpStatelessStreamableHttp|mcpStatelessStreamableHttp@io.ktor.server.application.Application(kotlin.String;kotlin.Boolean;kotlin.collections.List?;kotlin.collections.List?;kotlin.Function1){}[0] -final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcpStreamableHttp(kotlin/String = ..., kotlin/Boolean = ..., kotlin.collections/List? = ..., kotlin.collections/List? = ..., io.modelcontextprotocol.kotlin.sdk.server/EventStore? = ..., kotlin/Function1? = ..., kotlin/Function1) // io.modelcontextprotocol.kotlin.sdk.server/mcpStreamableHttp|mcpStreamableHttp@io.ktor.server.application.Application(kotlin.String;kotlin.Boolean;kotlin.collections.List?;kotlin.collections.List?;io.modelcontextprotocol.kotlin.sdk.server.EventStore?;kotlin.Function1?;kotlin.Function1){}[0] +final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcpStreamableHttp(kotlin/String = ..., kotlin/Boolean = ..., kotlin.collections/List? = ..., kotlin.collections/List? = ..., io.modelcontextprotocol.kotlin.sdk.server/EventStore? = ..., kotlin/Function1? = ..., kotlin.time/Duration = ..., kotlin/Int = ..., kotlin/Function1) // io.modelcontextprotocol.kotlin.sdk.server/mcpStreamableHttp|mcpStreamableHttp@io.ktor.server.application.Application(kotlin.String;kotlin.Boolean;kotlin.collections.List?;kotlin.collections.List?;io.modelcontextprotocol.kotlin.sdk.server.EventStore?;kotlin.Function1?;kotlin.time.Duration;kotlin.Int;kotlin.Function1){}[0] final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcpWebSocket(kotlin/Function0) // io.modelcontextprotocol.kotlin.sdk.server/mcpWebSocket|mcpWebSocket@io.ktor.server.application.Application(kotlin.Function0){}[0] final fun (io.ktor.server.application/Application).io.modelcontextprotocol.kotlin.sdk.server/mcpWebSocket(kotlin/String, kotlin/Function0) // io.modelcontextprotocol.kotlin.sdk.server/mcpWebSocket|mcpWebSocket@io.ktor.server.application.Application(kotlin.String;kotlin.Function0){}[0] final fun (io.ktor.server.routing/Route).io.modelcontextprotocol.kotlin.sdk.server/mcp(kotlin/Boolean = ..., kotlin.collections/List? = ..., kotlin.collections/List? = ..., kotlin/Long = ..., kotlin/Function1) // io.modelcontextprotocol.kotlin.sdk.server/mcp|mcp@io.ktor.server.routing.Route(kotlin.Boolean;kotlin.collections.List?;kotlin.collections.List?;kotlin.Long;kotlin.Function1){}[0] diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorServer.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorServer.kt index eba6f248d..f03daa4b7 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorServer.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorServer.kt @@ -6,10 +6,12 @@ import io.ktor.http.HttpMethod import io.ktor.http.HttpStatusCode import io.ktor.server.application.Application import io.ktor.server.application.ApplicationCall -import io.ktor.server.application.ApplicationCallPipeline import io.ktor.server.application.MissingApplicationPluginException +import io.ktor.server.application.createRouteScopedPlugin import io.ktor.server.application.install +import io.ktor.server.application.isHandled import io.ktor.server.application.pluginOrNull +import io.ktor.server.http.HttpRequestLifecycle import io.ktor.server.request.ApplicationRequest import io.ktor.server.request.header import io.ktor.server.request.httpMethod @@ -19,6 +21,7 @@ import io.ktor.server.routing.Route import io.ktor.server.routing.RoutingContext import io.ktor.server.routing.delete import io.ktor.server.routing.get +import io.ktor.server.routing.method import io.ktor.server.routing.post import io.ktor.server.routing.route import io.ktor.server.routing.routing @@ -29,7 +32,10 @@ import io.ktor.server.sse.heartbeat import io.ktor.server.sse.sse import io.ktor.utils.io.KtorDsl import io.modelcontextprotocol.kotlin.sdk.types.RPCError -import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineName +import kotlinx.coroutines.launch +import kotlin.time.Duration private val logger = KotlinLogging.logger {} @@ -40,6 +46,10 @@ private val logger = KotlinLogging.logger {} * **Precondition:** the [SSE] plugin must be installed on the application before calling this function. * Use [Application.mcp] if you want SSE to be installed automatically. * + * Installs [HttpRequestLifecycle] with `cancelCallOnClose` on the SSE `GET` route only, unless it is already + * installed on this route, a parent route, or the application, so that an SSE session ends when its client + * disconnects. + * * @param path the URL path to register the SSE endpoint. * @param enableDnsRebindingProtection whether to install [DnsRebindingProtection] on this route. Defaults to `true`. * @param allowedHosts hostnames allowed in the `Host` header. Defaults to `localhost`, `127.0.0.1`, `[::1]`. @@ -76,6 +86,10 @@ public fun Route.mcp( * **Precondition:** the [SSE] plugin must be installed on the application before calling this function. * Use [Application.mcp] if you want SSE to be installed automatically. * + * Installs [HttpRequestLifecycle] with `cancelCallOnClose` on the SSE `GET` route only, unless it is already + * installed on this route, a parent route, or the application, so that an SSE session ends when its client + * disconnects. + * * @param enableDnsRebindingProtection whether to install [DnsRebindingProtection] on this route. Defaults to `true`. * @param allowedHosts hostnames allowed in the `Host` header. Defaults to `localhost`, `127.0.0.1`, `[::1]`. * @param allowedOrigins origins allowed in the `Origin` header, compared by hostname only @@ -113,8 +127,12 @@ public fun Route.mcp( val transportManager = TransportManager() - sse { - mcpSseEndpoint("", transportManager, maxRequestBodySize, block) + // Only the SSE GET: this route may be the routing root, whose other handlers must not be cancelled. + method(HttpMethod.Get) { + installCancelCallOnClose() + sse { + mcpSseEndpoint("", transportManager, maxRequestBodySize, block) + } } post { @@ -128,7 +146,8 @@ public fun Route.mcp( * and sets up routing with the provided configuration block. * * Automatically installs [ContentNegotiation][io.ktor.server.plugins.contentnegotiation.ContentNegotiation] - * with [McpJson][io.modelcontextprotocol.kotlin.sdk.types.McpJson] and [SSE]. + * with [McpJson][io.modelcontextprotocol.kotlin.sdk.types.McpJson] and [SSE], and, like [Route.mcp], + * [HttpRequestLifecycle] on the SSE `GET` route only. * An already installed [SSE] plugin is reused. * * @param enableDnsRebindingProtection whether to install [DnsRebindingProtection] on this route. Defaults to `true`. @@ -170,6 +189,7 @@ private fun Application.mcpStreamableHttp( allowedOrigins: List?, configuration: StreamableHttpServerTransport.Configuration, sseHeartbeatConfig: (Heartbeat.() -> Unit)?, + sessions: StreamableHttpSessionManager, block: RoutingContext.() -> Server, ) { installMcpContentNegotiation() @@ -177,41 +197,34 @@ private fun Application.mcpStreamableHttp( install(SSE) } - val transportManager = TransportManager() + if (sessions.expires) { + // The application scope is cancelled when the application stops, which ends this loop. + launch(CoroutineName("mcp-session-expiry")) { sessions.closeExpiredSessionsPeriodically() } + } routing { route(path) { installDnsRebindingProtection(enableDnsRebindingProtection, allowedHosts, allowedOrigins) - - // Set Mcp-Session-Id on GET responses before Ktor's sse {} commits headers. - intercept(ApplicationCallPipeline.Plugins) { - if (context.request.httpMethod == HttpMethod.Get) { - val sessionId = context.request.header(MCP_SESSION_ID_HEADER) - if (sessionId != null && transportManager.getTransport(sessionId) != null) { - context.response.header(MCP_SESSION_ID_HEADER, sessionId) - } - } - } + installCancelCallOnClose() + installGetSessionGuard(sessions) sse { - val transport = existingStreamableTransport(call, transportManager) ?: return@sse + // Already checked by the session guard. + val transport = call.request.sessionId()?.let { sessions.touch(it) } ?: return@sse sseHeartbeatConfig?.let { config -> heartbeat(config) } transport.handleRequest(this, call) } post { - val transport = streamableTransport( - transportManager = transportManager, - configuration = configuration, - block = block, - ) ?: return@post - - transport.handleRequest(null, call) + if (call.request.sessionId() == null) { + openStreamableSession(sessions, configuration, block) + } else { + call.withStreamableSession(sessions) { transport -> transport.handleRequest(null, call) } + } } delete { - val transport = existingStreamableTransport(call, transportManager) ?: return@delete - transport.handleRequest(null, call) + call.withStreamableSession(sessions) { transport -> transport.handleRequest(null, call) } } } } @@ -228,6 +241,12 @@ private fun Application.mcpStreamableHttp( * with [McpJson][io.modelcontextprotocol.kotlin.sdk.types.McpJson] and [SSE]. * An already installed [SSE] plugin is reused. * + * A session expires after [sessionIdleTimeout] without requests (an open GET stream does not count), after which + * its id answers `404 Not Found`. At most [maxSessions] sessions are open at once; beyond that, a new session is + * refused with `503 Service Unavailable`. + * + * Also installs [HttpRequestLifecycle] with `cancelCallOnClose` on the route, unless it is already installed. + * * @param path The base path for the MCP Streamable HTTP endpoint. Defaults to "/mcp". * @param enableDnsRebindingProtection Enables DNS rebinding attack protection for the endpoint. Defaults to `true`. * @param allowedHosts A list of hostnames allowed to access the endpoint. @@ -240,8 +259,14 @@ private fun Application.mcpStreamableHttp( * @param eventStore An optional [EventStore] instance to enable resumable event stream functionality. * Allows storing and replaying events. * @param sseHeartbeatConfig The heartbeat configuration option for SSE connections. `null` means no heartbeat is sent. + * @param sessionIdleTimeout How long a session may go without activity before it is closed. Defaults to 30 minutes. + * Expiry is checked every few seconds, so a session can outlive this by up to 5 seconds. + * [Duration.INFINITE] disables expiry, so abandoned sessions pile up until [maxSessions] is reached. + * @param maxSessions The maximum number of sessions open at once, counting those still initializing. + * Defaults to 10,000. * @param block factory block with access to the [RoutingContext] (for reading request headers) * that creates and returns the [Server] to handle the connection. + * @throws IllegalArgumentException if [sessionIdleTimeout] or [maxSessions] is not positive. */ @KtorDsl public fun Application.mcpStreamableHttp( @@ -251,6 +276,8 @@ public fun Application.mcpStreamableHttp( allowedOrigins: List? = null, eventStore: EventStore? = null, sseHeartbeatConfig: (Heartbeat.() -> Unit)? = null, + sessionIdleTimeout: Duration = DEFAULT_SESSION_IDLE_TIMEOUT, + maxSessions: Int = DEFAULT_MAX_SESSIONS, block: RoutingContext.() -> Server, ) { mcpStreamableHttp( @@ -263,6 +290,7 @@ public fun Application.mcpStreamableHttp( enableJsonResponse = true, ), sseHeartbeatConfig = sseHeartbeatConfig, + sessions = StreamableHttpSessionManager(idleTimeout = sessionIdleTimeout, maxSessions = maxSessions), block = block, ) } @@ -376,19 +404,24 @@ private suspend fun ServerSSESession.mcpSseEndpoint( block: ServerSSESession.() -> Server, ) { val transport = mcpSseTransport(postEndpoint, transportManager, maxRequestBodySize) - - val server = block() - - server.onClose { + val transportClosed = CompletableDeferred() + // Registered before the session starts, so a close right after it starts is not missed. + transport.onClose { logger.info { "Server connection closed for sessionId: ${transport.sessionId}" } - transportManager.removeTransport(transport.sessionId) + transportClosed.complete(Unit) } - server.createSession(transport) - - logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" } + try { + block().createSession(transport) + logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" } - awaitCancellation() + // The connection lasts until the server closes the session, or until the client goes away, which + // cancels this call (see installCancelCallOnClose). + transportClosed.await() + } finally { + // Without its SSE connection the transport can no longer answer, so stop routing POSTs to it. + transportManager.removeTransport(transport.sessionId) + } } private fun ServerSSESession.mcpSseTransport( @@ -458,64 +491,116 @@ private suspend fun ApplicationCall.rejectUnsupportedMethod() { private fun ApplicationRequest.sessionId(): String? = header(MCP_SESSION_ID_HEADER) -private suspend fun existingStreamableTransport( - call: ApplicationCall, - transportManager: TransportManager, -): StreamableHttpServerTransport? { - val sessionId = call.request.sessionId() +/** + * Runs [block] with the transport of the session named by the `Mcp-Session-Id` header, keeping that session from + * expiring until [block] returns. Rejects the call when the header is missing or names no open session. + */ +private suspend fun ApplicationCall.withStreamableSession( + sessions: StreamableHttpSessionManager, + block: suspend (StreamableHttpServerTransport) -> Unit, +) { + val sessionId = sessionIdOrReject() ?: return + if (!sessions.withSession(sessionId, block)) rejectSessionNotFound() +} + +/** Returns the `Mcp-Session-Id` header, or answers `400 Bad Request` and returns `null` when it is missing. */ +private suspend fun ApplicationCall.sessionIdOrReject(): String? { + val sessionId = request.sessionId() if (sessionId.isNullOrEmpty()) { - call.reject( + reject( HttpStatusCode.BadRequest, RPCError.ErrorCode.CONNECTION_CLOSED, "Bad Request: No valid session ID provided", ) return null } - - val transport = transportManager.getTransport(sessionId) - return if (transport == null) { - call.reject( - HttpStatusCode.NotFound, - RPCError.ErrorCode.CONNECTION_CLOSED, - "Session not found", - ) - null - } else { - transport - } + return sessionId } -private suspend fun RoutingContext.streamableTransport( - transportManager: TransportManager, +/** + * Serves a POST that carries no `Mcp-Session-Id` on a new session. The session is kept only if the POST + * initializes it; otherwise it is closed once the POST completes. + */ +private suspend fun RoutingContext.openStreamableSession( + sessions: StreamableHttpSessionManager, configuration: StreamableHttpServerTransport.Configuration, block: RoutingContext.() -> Server, -): StreamableHttpServerTransport? { - val sessionId = call.request.sessionId() - if (sessionId != null) { - val transport = transportManager.getTransport(sessionId) - return transport ?: existingStreamableTransport(call, transportManager) +) { + if (!sessions.tryReserve()) { + logger.warn { "Rejecting a new StreamableHttp session: ${sessions.maxSessions} sessions are already open" } + call.reject( + HttpStatusCode.ServiceUnavailable, + RPCError.ErrorCode.INTERNAL_ERROR, + "Service Unavailable: Too many open sessions", + ) + return } val transport = StreamableHttpServerTransport(configuration) transport.setOnSessionInitialized { initializedSessionId -> - transportManager.addTransport(initializedSessionId, transport) + sessions.register(initializedSessionId, transport) logger.info { "New StreamableHttp connection established and stored with sessionId: $initializedSessionId" } } transport.setOnSessionClosed { closedSession -> - transportManager.removeTransport(closedSession) + sessions.remove(closedSession, transport) logger.info { "Closed StreamableHttp connection and removed sessionId: $closedSession" } } - val server = block() - server.onClose { - transport.sessionId?.let { transportManager.removeTransport(it) } - logger.info { "Server connection closed for sessionId: ${transport.sessionId}" } + try { + val session = block().createSession(transport) + // Registered on the session rather than on the Server, which may be shared by every session. + session.onClose { + transport.sessionId?.let { sessionId -> + sessions.remove(sessionId, transport) + logger.debug { "Server connection closed for sessionId: $sessionId" } + } + } + transport.handleRequest(null, call) + } finally { + // The transport gets its session id when it initializes, which is also when the session is registered. + val sessionId = transport.sessionId + if (sessionId != null) { + sessions.release(sessionId, transport) + } else { + sessions.cancelReservation() + // No request can reach a transport that never initialized. Close it so that a Server shared + // across sessions drops the session created for it. + transport.close() + } } - server.createSession(transport) +} - return transport +/** + * Has the engine cancel a call once its client disconnects. Otherwise a handler that only waits, like the one + * holding an SSE connection open, keeps running after its client is gone, and so does its session. Only the CIO + * and Netty engines support this; others notice a gone client only when writing to it. An [HttpRequestLifecycle] + * already installed on this route, a parent route, or the application is kept as it is. + */ +private fun Route.installCancelCallOnClose() { + try { + plugin(HttpRequestLifecycle) + } catch (_: MissingApplicationPluginException) { + install(HttpRequestLifecycle) { cancelCallOnClose = true } + } +} + +/** Checks the session of a GET before sse {}, which sends 200 and the headers before its handler runs. */ +private fun Route.installGetSessionGuard(sessions: StreamableHttpSessionManager) { + val guard = createRouteScopedPlugin("MCP-StreamableHttpGetSessionGuard") { + onCall { call -> + if (call.request.httpMethod != HttpMethod.Get || call.isHandled) return@onCall + val sessionId = call.sessionIdOrReject() ?: return@onCall + // Touched so that the session cannot expire before sse {} runs. + if (sessions.touch(sessionId) == null) { + call.rejectSessionNotFound() + return@onCall + } + call.response.header(MCP_SESSION_ID_HEADER, sessionId) + } + } + install(guard) } private fun Route.installDnsRebindingProtection(enabled: Boolean, hosts: List?, origins: List?) { 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 e85d71bfa..b59d1f781 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 @@ -33,11 +33,11 @@ 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 import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeoutOrNull import kotlinx.serialization.json.JsonArray import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.decodeFromJsonElement @@ -46,6 +46,7 @@ import kotlin.concurrent.atomics.ExperimentalAtomicApi import kotlin.coroutines.cancellation.CancellationException import kotlin.time.Duration import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.seconds import kotlin.uuid.ExperimentalUuidApi import kotlin.uuid.Uuid @@ -54,12 +55,25 @@ private const val MCP_PROTOCOL_VERSION_HEADER = "mcp-protocol-version" private const val MCP_RESUMPTION_TOKEN_HEADER = "Last-Event-ID" private const val MIN_PRIMING_EVENT_PROTOCOL_VERSION = "2025-11-25" +/** How long closing an SSE session may wait on a slow client before giving up. */ +private val SSE_SESSION_CLOSE_TIMEOUT: Duration = 1.seconds + /** * A holder for an active request call. * If [StreamableHttpServerTransport.Configuration.enableJsonResponse] is true, the session is null. * Otherwise, the session is not null. */ -private data class SessionContext(val session: ServerSSESession?, val call: ApplicationCall) +private class SessionContext(val session: ServerSSESession?, val call: ApplicationCall) { + /** Completed when the transport drops this stream; the GET handler waits on it. */ + val released: CompletableDeferred = CompletableDeferred() +} + +/** Why a POST was not registered. */ +private sealed interface PostRefusal { + data object TransportClosed : PostRefusal + + class ReusedId(val id: RequestId) : PostRefusal +} /** * Server transport for Streamable HTTP: this implements the MCP Streamable HTTP transport specification. @@ -197,6 +211,9 @@ public class StreamableHttpServerTransport(private val configuration: Configurat private val sessionMutex = Mutex() private val streamMutex = Mutex() + /** Set by [close]; later requests are refused. */ + private val closed: AtomicBoolean = AtomicBoolean(false) + private companion object { const val STANDALONE_SSE_STREAM_ID = "_GET_stream" } @@ -313,26 +330,27 @@ public class StreamableHttpServerTransport(private val configuration: Configurat // (disconnect or cancellation), so the responses are undeliverable and dropped quietly. val pendingPost = pendingJsonResponses.remove(streamId) if (pendingPost == null) { - cleanUpRequests(relatedIds) + cleanUpRequests(relatedIds, streamId) return } val responses = relatedIds.mapNotNull { requestToResponseMapping[it] } - cleanUpRequests(relatedIds) + cleanUpRequests(relatedIds, streamId) pendingPost.complete(responses) } else { - // 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) + cleanUpRequests(relatedIds, streamId) } } + // Every response is out, so end the stream; outside the lock, as closing can wait on a slow client. + if (!configuration.enableJsonResponse) activeStream?.session?.closeQuietly() } - private fun cleanUpRequests(ids: Collection) { + /** Retires [ids] that are still routed to [streamId], leaving any later request that reused an id alone. */ + private fun cleanUpRequests(ids: Collection, streamId: String) { ids.forEach { requestId -> - requestToResponseMapping.remove(requestId) - requestToStreamMapping.remove(requestId) - cancelledRequestIds.remove(requestId) + if (requestToStreamMapping.remove(requestId, streamId)) { + requestToResponseMapping.remove(requestId) + cancelledRequestIds.remove(requestId) + } } } @@ -369,20 +387,16 @@ public class StreamableHttpServerTransport(private val configuration: Configurat if (!allSettled) return val pendingPost = pendingJsonResponses.remove(streamId) ?: return val responses = relatedIds.mapNotNull { requestToResponseMapping[it] } - cleanUpRequests(relatedIds) + cleanUpRequests(relatedIds, streamId) pendingPost.complete(responses) } } override suspend fun close() { withContext(NonCancellable) { - streamMutex.withLock { - streamsMapping.values.forEach { - try { - it.session?.close() - } catch (_: Exception) { - } - } + val streams = streamMutex.withLock { + closed.store(true) + val contexts = streamsMapping.values.toList() streamsMapping.clear() requestToStreamMapping.clear() requestToResponseMapping.clear() @@ -390,8 +404,13 @@ public class StreamableHttpServerTransport(private val configuration: Configurat // Release any POST calls suspended awaiting JSON responses. pendingJsonResponses.values.forEach { it.cancel() } pendingJsonResponses.clear() - invokeOnCloseCallback() + contexts } + // Closing an SSE session can wait on a slow client, so it happens outside the lock, bounded, and after + // the close callbacks, which cancel the server's own pending sends. + streams.forEach { it.released.complete(Unit) } + invokeOnCloseCallback() + streams.forEach { it.session?.closeQuietly() } } } @@ -532,10 +551,33 @@ public class StreamableHttpServerTransport(private val configuration: Configurat null } - streamMutex.withLock { + val requestIds = messages.filterIsInstance().map { it.id } + val refusal = streamMutex.withLock { + // close() may have run since validateSession. + if (closed.load()) return@withLock PostRefusal.TransportClosed + // MCP forbids reusing a request id in a session: a reused in-flight id would take its response. + val seen = HashSet() + val reusedId = requestIds.firstOrNull { !seen.add(it) || it in requestToStreamMapping } + if (reusedId != null) return@withLock PostRefusal.ReusedId(reusedId) streamsMapping[streamId] = SessionContext(session, call) - messages.filterIsInstance().forEach { requestToStreamMapping[it.id] = streamId } + requestIds.forEach { requestToStreamMapping[it] = streamId } pendingPost?.let { pendingJsonResponses[streamId] = it } + null + } + if (refusal != null) { + // An SSE response has already started, so only a JSON one can still carry the error. + if (pendingPost == null) return + when (refusal) { + PostRefusal.TransportClosed -> call.rejectSessionNotFound(requestIds.firstOrNull()) + + is PostRefusal.ReusedId -> call.reject( + HttpStatusCode.BadRequest, + RPCError.ErrorCode.INVALID_REQUEST, + "Invalid Request: Request id is already in use", + refusal.id, + ) + } + return } call.coroutineContext.job.invokeOnCompletion { streamsMapping.remove(streamId) } @@ -568,7 +610,7 @@ public class StreamableHttpServerTransport(private val configuration: Configurat // the success path (send() or retirement already cleaned up). pendingJsonResponses.remove(streamId) pendingPost.cancel() - cleanUpRequests(messages.filterIsInstance().map { it.id }) + cleanUpRequests(requestIds, streamId) } } catch (e: CancellationException) { throw e @@ -612,32 +654,16 @@ public class StreamableHttpServerTransport(private val configuration: Configurat flushSse(sseSession) val newContext = SessionContext(sseSession, call) streamMutex.withLock { - streamsMapping[STANDALONE_SSE_STREAM_ID]?.let { existingContext -> - // Close the previous SSE session. If alive, this cancels the old - // coroutine (which will hit its identity-guarded finally — that finally - // won't double-remove, since we replace the mapping below). - try { - existingContext.session?.close() - } catch (e: CancellationException) { - throw e - } catch (_: Exception) { - // Ignore — the old stream may already be closed. - } - // Evict the stale mapping — the old session is closed either way. - streamsMapping.remove(STANDALONE_SSE_STREAM_ID) - } - streamsMapping[STANDALONE_SSE_STREAM_ID] = newContext + if (closed.load()) return + // Release the stream this one replaces; its own handler then returns and closes it. + streamsMapping.put(STANDALONE_SSE_STREAM_ID, newContext)?.released?.complete(Unit) } val clientProtocolVersion = call.request.header(MCP_PROTOCOL_VERSION_HEADER) ?: DEFAULT_NEGOTIATED_PROTOCOL_VERSION maybeSendPrimingEvent(STANDALONE_SSE_STREAM_ID, sseSession, clientProtocolVersion) - // Keep the SSE connection open until the client disconnects or the transport is closed. - // Cleanup uses try/finally (runs during cancellation propagation) instead of - // invokeOnCompletion (runs after job completion) to minimize the window between - // disconnect and mapping removal. Identity check ensures only this stream's entry - // is removed — not a replacement that arrived in the meantime. + // Keep the stream open until the client leaves, a newer GET replaces it, or the transport closes. try { - awaitCancellation() + newContext.released.await() } finally { // Atomic compare-and-remove — only evict our own entry, not a replacement's. streamsMapping.remove(STANDALONE_SSE_STREAM_ID, newContext) @@ -765,19 +791,12 @@ public class StreamableHttpServerTransport(private val configuration: Configurat return false } - return when (headerId) { - sessionId -> true - - else -> { - call.reject( - HttpStatusCode.NotFound, - REQUEST_TIMEOUT, - "Session not found", - id, - ) - false - } + // A closed transport's session has ended, and the spec requires 404 for it. + if (headerId != sessionId || closed.load()) { + call.rejectSessionNotFound(id) + return false } + return true } private suspend fun validateProtocolVersion(call: ApplicationCall, id: RequestId? = null): Boolean { @@ -892,17 +911,20 @@ public class StreamableHttpServerTransport(private val configuration: Configurat try { session?.send(event = "message", id = eventId, data = McpJson.encodeToString(message)) } catch (e: CancellationException) { - // Atomic compare-and-remove — only evict this stream's entry, not a replacement's. - val current = streamsMapping[streamId] - if (current != null && current.session === session) { - streamsMapping.remove(streamId, current) - } + releaseStream(streamId, session) throw e } catch (_: Exception) { - val current = streamsMapping[streamId] - if (current != null && current.session === session) { - streamsMapping.remove(streamId, current) - } + releaseStream(streamId, session) + } + } + + /** Evicts the stream [streamId] if [session] still serves it, and releases a GET handler waiting on it. */ + private fun releaseStream(streamId: String, session: ServerSSESession?) { + // Atomic compare-and-remove — only evict this stream's entry, not a replacement's. + val current = streamsMapping[streamId] + if (current != null && current.session === session) { + streamsMapping.remove(streamId, current) + current.released.complete(Unit) } } @@ -950,3 +972,19 @@ internal suspend fun ApplicationCall.reject( val error = JSONRPCError(id = id, error = RPCError(code = code, message = message)) respondText(McpJson.encodeToString(error), ContentType.Application.Json, status) } + +/** Answers a request for a session that does not exist or has ended. */ +internal suspend fun ApplicationCall.rejectSessionNotFound(id: RequestId? = null) { + reject(HttpStatusCode.NotFound, REQUEST_TIMEOUT, "Session not found", id) +} + +/** Closes this session, ignoring failures and giving up after [SSE_SESSION_CLOSE_TIMEOUT]. */ +private suspend fun ServerSSESession.closeQuietly() { + try { + withTimeoutOrNull(SSE_SESSION_CLOSE_TIMEOUT) { close() } + } catch (e: CancellationException) { + throw e + } catch (_: Exception) { + // Ignore — the stream may already be closed. + } +} diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionManager.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionManager.kt new file mode 100644 index 000000000..47f56ad29 --- /dev/null +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionManager.kt @@ -0,0 +1,226 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.github.oshai.kotlinlogging.KotlinLogging +import io.modelcontextprotocol.kotlin.sdk.shared.Transport +import kotlinx.atomicfu.AtomicRef +import kotlinx.atomicfu.atomic +import kotlinx.atomicfu.loop +import kotlinx.atomicfu.update +import kotlinx.collections.immutable.PersistentMap +import kotlinx.collections.immutable.persistentMapOf +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import kotlin.coroutines.cancellation.CancellationException +import kotlin.time.Duration +import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.minutes +import kotlin.time.Duration.Companion.seconds +import kotlin.time.TimeMark +import kotlin.time.TimeSource + +private val logger = KotlinLogging.logger {} + +/** How long a stateful Streamable HTTP session may stay idle before [mcpStreamableHttp] closes it. */ +internal val DEFAULT_SESSION_IDLE_TIMEOUT: Duration = 30.minutes + +/** How many stateful Streamable HTTP sessions [mcpStreamableHttp] keeps open at once. */ +internal const val DEFAULT_MAX_SESSIONS: Int = 10_000 + +private val MIN_EXPIRY_CHECK_INTERVAL: Duration = 10.milliseconds +private val MAX_EXPIRY_CHECK_INTERVAL: Duration = 5.seconds + +/** + * Tracks the sessions of a stateful [mcpStreamableHttp] endpoint and bounds how many it keeps and for how long. + * + * A session is idle while none of the requests run through [withSession] is in flight. Once it has been idle + * for [idleTimeout], [closeExpiredSessions] forgets it, so its id answers `404`, and closes its transport, which + * closes the server session and any GET stream on it too. At most [maxSessions] sessions are open at once, + * counting the ones still initializing. + * + * @param T the transport a session is served by, a [StreamableHttpServerTransport] outside tests + */ +internal class StreamableHttpSessionManager( + private val idleTimeout: Duration = DEFAULT_SESSION_IDLE_TIMEOUT, + val maxSessions: Int = DEFAULT_MAX_SESSIONS, + private val timeSource: TimeSource = TimeSource.Monotonic, +) { + init { + require(idleTimeout.isPositive()) { "sessionIdleTimeout must be positive, but was $idleTimeout" } + require(maxSessions > 0) { "maxSessions must be positive, but was $maxSessions" } + } + + private val sessions: AtomicRef> = atomic(persistentMapOf()) + + /** Registered sessions plus reservations for sessions that are still initializing. */ + private val openCount = atomic(0) + + /** Whether idle sessions expire at all; `false` when [idleTimeout] is infinite. */ + val expires: Boolean get() = idleTimeout.isFinite() + + /** How often [closeExpiredSessionsPeriodically] looks for expired sessions. */ + val expiryCheckInterval: Duration = idleTimeout.coerceIn(MIN_EXPIRY_CHECK_INTERVAL, MAX_EXPIRY_CHECK_INTERVAL) + + operator fun contains(sessionId: String): Boolean = sessionId in sessions.value + + /** Claims room for a new session, or returns `false` when [maxSessions] sessions are already open. */ + fun tryReserve(): Boolean { + openCount.loop { count -> + if (count >= maxSessions) return false + if (openCount.compareAndSet(count, count + 1)) return true + } + } + + /** Gives back a reservation whose session never initialized. */ + fun cancelReservation() { + openCount.decrementAndGet() + } + + /** + * Turns a reservation into a session reachable as [sessionId]. The session starts with its initializing + * request in flight, which the caller must [release]. + */ + fun register(sessionId: String, transport: T) { + val session = Session(sessionId, transport) + sessions.update { it.putting(sessionId, session) } + } + + /** Ends a request on the session that [transport] serves as [sessionId]. Does nothing once it is gone. */ + fun release(sessionId: String, transport: T) { + find(sessionId, transport)?.release() + } + + /** + * Runs [block] with the transport of [sessionId] and keeps the session from expiring until it returns. + * Returns `false` without running [block] when no such session is open. + */ + suspend fun withSession(sessionId: String, block: suspend (T) -> Unit): Boolean { + val session = sessions.value[sessionId] + if (session == null || !session.tryAcquire()) return false + try { + block(session.transport) + } finally { + session.release() + } + return true + } + + /** + * Returns the transport of [sessionId] and restarts its idle timer, or `null` when no such session is open. + * + * Unlike [withSession], this does not keep the session from expiring afterwards. It serves GET streams: + * the engine may never notice that the client of a stream went away, so a stream must not keep its + * session open. Expiry closes the stream along with the session. + */ + fun touch(sessionId: String): T? { + val session = sessions.value[sessionId] ?: return null + return if (session.touch()) session.transport else null + } + + /** Forgets the session that [transport] serves as [sessionId]. Does nothing once it is gone. */ + fun remove(sessionId: String, transport: T) { + find(sessionId, transport)?.let { forget(it) } + } + + /** Closes every session that has been idle for [idleTimeout], each on its own so that none waits on another. */ + suspend fun closeExpiredSessions(): Unit = coroutineScope { + forgetExpiredSessions().forEach { session -> launch { close(session) } } + } + + /** + * Closes expired sessions every [expiryCheckInterval] until cancelled. Each sweep runs on its own, so a close + * that hangs, for instance on a client that stopped reading its stream, cannot hold up the next ones. + */ + suspend fun closeExpiredSessionsPeriodically(): Unit = coroutineScope { + while (true) { + delay(expiryCheckInterval) + launch { + // A failed sweep must not end expiry for good; the next one tries again. + try { + closeExpiredSessions() + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + logger.error(e) { "Failed to expire idle StreamableHttp sessions" } + } + } + } + } + + private fun find(sessionId: String, transport: T): Session? = + sessions.value[sessionId]?.takeIf { it.transport === transport } + + /** Forgets every session that has been idle for [idleTimeout], so its id answers `404`, and returns them. */ + private fun forgetExpiredSessions(): List = buildList { + for (session in sessions.value.values) { + if (session.tryExpire() && forget(session)) { + logger.info { "Closing StreamableHttp session ${session.id} after $idleTimeout without activity" } + add(session) + } + } + } + + /** Stops routing requests to [session] and frees its slot. Returns `false` if it was already forgotten. */ + private fun forget(session: Session): Boolean { + session.markClosed() + sessions.loop { current -> + if (current[session.id] !== session) return false + if (sessions.compareAndSet(current, current.removing(session.id))) { + openCount.decrementAndGet() + return true + } + } + } + + private suspend fun close(session: Session) { + try { + session.transport.close() + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + logger.warn(e) { "Failed to close expired StreamableHttp session ${session.id}" } + } + } + + /** A registered session: its transport plus the requests in flight for it. */ + private inner class Session(val id: String, val transport: T) { + // A single immutable snapshot so that expiry and a starting request race on one compare-and-set. + private val activity = atomic(Activity(inFlight = 1, lastActive = timeSource.markNow(), closed = false)) + + fun tryAcquire(): Boolean { + activity.loop { current -> + if (current.closed) return false + if (activity.compareAndSet(current, current.copy(inFlight = current.inFlight + 1))) return true + } + } + + fun release() { + val now = timeSource.markNow() + activity.update { it.copy(inFlight = it.inFlight - 1, lastActive = now) } + } + + fun touch(): Boolean { + val now = timeSource.markNow() + activity.loop { current -> + if (current.closed) return false + if (activity.compareAndSet(current, current.copy(lastActive = now))) return true + } + } + + /** Marks the session closed if it has been idle for [idleTimeout]. */ + fun tryExpire(): Boolean { + activity.loop { current -> + if (current.closed || current.inFlight > 0 || current.lastActive.elapsedNow() < idleTimeout) { + return false + } + if (activity.compareAndSet(current, current.copy(closed = true))) return true + } + } + + fun markClosed() { + activity.update { it.copy(closed = true) } + } + } + + private data class Activity(val inFlight: Int, val lastActive: TimeMark, val closed: Boolean) +} diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/DnsRebindingProtectionTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/DnsRebindingProtectionTest.kt index 191a65ea0..c3969b54d 100644 --- a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/DnsRebindingProtectionTest.kt +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/DnsRebindingProtectionTest.kt @@ -6,6 +6,7 @@ import io.kotest.matchers.shouldBe import io.kotest.matchers.shouldNotBe import io.kotest.matchers.string.shouldContain import io.kotest.matchers.string.shouldNotContain +import io.ktor.client.request.get import io.ktor.client.request.header import io.ktor.client.request.post import io.ktor.client.statement.bodyAsText @@ -247,6 +248,18 @@ class DnsRebindingProtectionTest { response.shouldHaveStatus(HttpStatusCode.Forbidden) } + @Test + fun `mcpStreamableHttp rejects a GET from a non-localhost Host before checking its session`() = testApplication { + application { + mcpStreamableHttp { testServer() } + } + + val response = client.get("/mcp") { + header(HttpHeaders.Host, "evil.com") + } + response.shouldHaveStatus(HttpStatusCode.Forbidden) + } + @Test fun `mcpStatelessStreamableHttp rejects non-localhost Host by default`() = testApplication { application { diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorExtensionsTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorExtensionsTest.kt index 69410c5ea..a132341ad 100644 --- a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorExtensionsTest.kt +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorExtensionsTest.kt @@ -2,7 +2,9 @@ package io.modelcontextprotocol.kotlin.sdk.server import io.kotest.assertions.ktor.client.shouldHaveContentType import io.kotest.assertions.ktor.client.shouldHaveStatus +import io.kotest.assertions.nondeterministic.eventually import io.kotest.assertions.throwables.shouldThrow +import io.kotest.matchers.maps.shouldHaveSize import io.kotest.matchers.nulls.shouldNotBeNull import io.kotest.matchers.string.shouldContain import io.ktor.client.HttpClient @@ -12,6 +14,7 @@ import io.ktor.client.request.header import io.ktor.client.request.post import io.ktor.client.request.prepareGet import io.ktor.client.request.setBody +import io.ktor.client.statement.HttpResponse import io.ktor.http.ContentType import io.ktor.http.HttpHeaders import io.ktor.http.HttpStatusCode @@ -24,6 +27,7 @@ import io.ktor.server.testing.testApplication import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.McpJson import kotlin.test.Test +import kotlin.time.Duration.Companion.seconds class KtorExtensionsTest { @@ -71,6 +75,32 @@ class KtorExtensionsTest { client.post("/api").shouldHaveStatus(HttpStatusCode.NotFound) } + /** + * Verifies that a closed session stops receiving POSTs and is not retained until the + * [Server] itself closes, even while its SSE connection is still open. + */ + @Test + fun `Route mcp stops routing POSTs to a session once it closes`() = testApplication { + val server = testServer() + application { + install(SSE) + routing { + mcp(enableDnsRebindingProtection = false) { server } + } + } + + client.prepareGet("/").execute { response -> + val sessionId = response.readSseSessionId().shouldNotBeNull() + client.postPing("/", sessionId).shouldHaveStatus(HttpStatusCode.Accepted) + // The session registers with the Server just after the endpoint event goes out. + eventually(5.seconds) { server.sessions shouldHaveSize 1 } + + server.sessions.values.single().close() + + client.postPing("/", sessionId).shouldHaveStatus(HttpStatusCode.NotFound) + } + } + @Test fun `Application mcp should install SSE and register endpoints at the root`() = testApplication { application { @@ -128,12 +158,16 @@ class KtorExtensionsTest { response.shouldHaveContentType(ContentType.Text.EventStream) val sessionId = response.readSseSessionId().shouldNotBeNull() - post("$path?sessionId=$sessionId") { - contentType(ContentType.Application.Json) - setBody("""{"jsonrpc":"2.0","id":1,"method":"ping"}""") - }.shouldHaveStatus(HttpStatusCode.Accepted) + postPing(path, sessionId).shouldHaveStatus(HttpStatusCode.Accepted) } post(path).shouldHaveStatus(HttpStatusCode.BadRequest) } + + /** POSTs a JSON-RPC ping to the session [sessionId] of the SSE endpoint at [path]. */ + private suspend fun HttpClient.postPing(path: String, sessionId: String): HttpResponse = + post("$path?sessionId=$sessionId") { + contentType(ContentType.Application.Json) + setBody("""{"jsonrpc":"2.0","id":1,"method":"ping"}""") + } } diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpResponseDeliveryTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpResponseDeliveryTest.kt index dbcf610f3..5819f710c 100644 --- a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpResponseDeliveryTest.kt +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpResponseDeliveryTest.kt @@ -2,6 +2,7 @@ package io.modelcontextprotocol.kotlin.sdk.server import io.kotest.matchers.shouldBe import io.ktor.client.request.delete +import io.ktor.client.request.get import io.ktor.client.request.header import io.ktor.client.request.request import io.ktor.client.statement.HttpResponse @@ -11,6 +12,9 @@ import io.ktor.http.HttpHeaders import io.ktor.http.HttpMethod import io.ktor.http.HttpStatusCode import io.ktor.http.contentType +import io.ktor.server.response.respondText +import io.ktor.server.routing.get +import io.ktor.server.routing.routing import io.ktor.server.testing.testApplication import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError import io.modelcontextprotocol.kotlin.sdk.types.McpJson @@ -59,5 +63,46 @@ class StreamableHttpResponseDeliveryTest { response.decodeError().error.message shouldBe "Session not found" } + @Test + fun `stateful GET for an unknown session returns 404`() = testApplication { + application { mcpStreamableHttp { testServer() } } + + val response = client.get("/mcp") { + header(HttpHeaders.Host, "localhost") + header(HttpHeaders.Accept, eventStream) + header(MCP_SESSION_ID_HEADER, "unknown-session") + } + + response.status shouldBe HttpStatusCode.NotFound + val error = response.decodeError().error + error.message shouldBe "Session not found" + error.code shouldBe -32001 + } + + @Test + fun `stateful GET without a session id returns 400`() = testApplication { + application { mcpStreamableHttp { testServer() } } + + val response = client.get("/mcp") { + header(HttpHeaders.Host, "localhost") + header(HttpHeaders.Accept, eventStream) + } + + response.status shouldBe HttpStatusCode.BadRequest + } + + @Test + fun `GET routes next to the stateful endpoint are not checked for a session`() = testApplication { + application { + mcpStreamableHttp { testServer() } + routing { get("/health") { call.respondText("ok") } } + } + + val response = client.get("/health") + + response.status shouldBe HttpStatusCode.OK + response.bodyAsText() shouldBe "ok" + } + private suspend fun HttpResponse.decodeError(): JSONRPCError = McpJson.decodeFromString(bodyAsText()) } 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 a5d299ce2..0b7b3cdbc 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 @@ -21,12 +21,15 @@ 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.get 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.discard import io.ktor.utils.io.readLine @@ -53,6 +56,7 @@ import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.async +import kotlinx.coroutines.awaitCancellation import kotlinx.coroutines.cancel import kotlinx.coroutines.launch import kotlinx.coroutines.test.runTest @@ -62,6 +66,7 @@ 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.time.Duration.Companion.seconds import io.ktor.client.plugins.contentnegotiation.ContentNegotiation as ClientContentNegotiation @@ -439,6 +444,33 @@ class StreamableHttpServerTransportTest { } } + @Test + fun `sse terminal response is not held up by a stream that never closes`() = testApplication { + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = false), + ) + transport.setSessionIdGenerator(null) // stateless: accept requests without an init handshake + transport.answerRequests() + val handled = CompletableDeferred() + application { + routing { + post(path) { + transport.handlePostRequest(UnclosableServerSSESession(call, call.coroutineContext), call) + handled.complete(Unit) + } + } + } + + val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + try { + scope.launch { client.postMessages(null, JSONRPCRequest(id = RequestId("ping-1"), method = "ping")) } + // The response is sent, but closing its stream never finishes; send() must return anyway. + withTimeout(5.seconds) { handled.await() } + } finally { + scope.cancel() + } + } + @ParameterizedTest @MethodSource("unparsableBodyCases") fun `POST with an unparsable body returns a JSON-RPC error`(body: String, expectedCode: Int) = testApplication { @@ -615,3 +647,14 @@ class StreamableHttpServerTransportTest { private fun toolsList(id: RequestId): JSONRPCRequest = JSONRPCRequest(id = id, method = Method.Defined.ToolsList.value) } + +/** An SSE session whose client stopped reading: closing it waits on a stuck send and never finishes. */ +private class UnclosableServerSSESession( + override val call: ApplicationCall, + override val coroutineContext: CoroutineContext, +) : ServerSSESession { + override suspend fun send(event: ServerSentEvent) {} + override suspend fun close() { + awaitCancellation() + } +} diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionLimitsTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionLimitsTest.kt new file mode 100644 index 000000000..b738b4b0d --- /dev/null +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionLimitsTest.kt @@ -0,0 +1,263 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.kotest.assertions.nondeterministic.eventually +import io.kotest.matchers.maps.shouldBeEmpty +import io.kotest.matchers.maps.shouldHaveSize +import io.kotest.matchers.nulls.shouldNotBeNull +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldContain +import io.ktor.client.HttpClient +import io.ktor.client.request.HttpRequestBuilder +import io.ktor.client.request.delete +import io.ktor.client.request.get +import io.ktor.client.request.header +import io.ktor.client.request.post +import io.ktor.client.request.setBody +import io.ktor.client.statement.HttpResponse +import io.ktor.client.statement.bodyAsText +import io.ktor.http.HttpHeaders +import io.ktor.http.HttpStatusCode +import io.ktor.server.testing.testApplication +import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import io.modelcontextprotocol.kotlin.sdk.types.RPCError +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.async +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.delay +import kotlinx.coroutines.withTimeout +import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.Test +import kotlin.test.assertFailsWith +import kotlin.time.Duration +import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.seconds + +/** + * Verifies that a stateful Streamable HTTP endpoint bounds the sessions it keeps: idle sessions expire, + * at most `maxSessions` are open at once, every way a session ends gives its slot back, and a client cannot + * keep a session from expiring by reusing a request id. + */ +class StreamableHttpSessionLimitsTest { + + private val idleTimeout = 200.milliseconds + + @Test + fun `abandoned session is closed once idle for sessionIdleTimeout`() = testApplication { + val server = testServer() + application { mcpStreamableHttp(sessionIdleTimeout = idleTimeout) { server } } + + val sessionId = client.initialize() + + eventually(5.seconds) { + server.sessions.shouldBeEmpty() + } + client.ping(sessionId).status shouldBe HttpStatusCode.NotFound + } + + @Test + fun `requests keep a session alive`() = testApplication { + val server = testServer() + application { mcpStreamableHttp(sessionIdleTimeout = 400.milliseconds) { server } } + + val sessionId = client.initialize() + // Well past the timeout plus two expiry checks, so a session that requests did not keep alive expires. + repeat(13) { + delay(100.milliseconds) + client.ping(sessionId).status shouldBe HttpStatusCode.OK + } + server.sessions shouldHaveSize 1 + } + + @Test + fun `session beyond maxSessions is rejected with 503 until one closes`() = testApplication { + application { mcpStreamableHttp(maxSessions = 2) { testServer() } } + + val first = client.initialize() + client.initialize() + + val rejected = client.post("/mcp") { + localhostStreamableHeaders() + setBody(initializeRequestBody()) + } + rejected.status shouldBe HttpStatusCode.ServiceUnavailable + rejected.headers[MCP_SESSION_ID_HEADER] shouldBe null + rejected.decodeError().error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR + + client.delete("/mcp") { + header(HttpHeaders.Host, "localhost") + header(MCP_SESSION_ID_HEADER, first) + }.status shouldBe HttpStatusCode.OK + + client.initialize() + } + + @Test + fun `a session the server closes frees its slot and its id answers 404`() = testApplication { + val server = testServer() + application { mcpStreamableHttp(maxSessions = 1) { server } } + + val sessionId = client.initialize() + server.sessions.values.single().close() + + client.ping(sessionId).status shouldBe HttpStatusCode.NotFound + client.initialize() + } + + @Test + fun `a failing server factory gives its session slot back`() = testApplication { + val factoryCalls = AtomicInteger() + application { + mcpStreamableHttp(maxSessions = 1) { + check(factoryCalls.getAndIncrement() > 0) { "The first server fails to start" } + testServer() + } + } + + client.post("/mcp") { + localhostStreamableHeaders() + setBody(initializeRequestBody()) + }.status shouldBe HttpStatusCode.InternalServerError + client.initialize() + } + + @Test + fun `requests that never initialize leave no session behind on a shared server`() = testApplication { + val server = testServer() + application { mcpStreamableHttp(maxSessions = 1) { server } } + + repeat(3) { + client.post("/mcp") { + localhostStreamableHeaders() + setBody("""{"jsonrpc":"2.0","id":1,"method":"tools/list"}""") + }.status shouldBe HttpStatusCode.BadRequest + } + + eventually(5.seconds) { + server.sessions.shouldBeEmpty() + } + // The rejected requests gave their slot back. + client.initialize() + server.sessions shouldHaveSize 1 + } + + @Test + fun `a request id still in flight cannot be reused`() = testApplication { + val toolStarted = CompletableDeferred() + val releaseTool = CompletableDeferred() + application { + mcpStreamableHttp { + Server( + Implementation("test-server", "1.0"), + ServerOptions(capabilities = ServerCapabilities(tools = ServerCapabilities.Tools())), + ) { + addTool(name = "wait", description = "Answers once the test releases it") { + toolStarted.complete(Unit) + releaseTool.await() + CallToolResult(content = listOf(TextContent("released"))) + } + } + } + } + val sessionId = client.initialize() + client.sendInitialized(sessionId) + + coroutineScope { + val inFlight = async { + client.post("/mcp") { + localhostStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody( + """{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"wait","arguments":{}}}""", + ) + } + } + toolStarted.await() + + // Accepting the reused id would route the answer for id 7 to this POST and leave the first one + // waiting forever, holding its session open. + val reused = client.post("/mcp") { + localhostStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody("""{"jsonrpc":"2.0","id":7,"method":"ping"}""") + } + reused.status shouldBe HttpStatusCode.BadRequest + reused.decodeError().error.code shouldBe RPCError.ErrorCode.INVALID_REQUEST + + releaseTool.complete(Unit) + val answered = withTimeout(5.seconds) { inFlight.await() } + answered.status shouldBe HttpStatusCode.OK + answered.bodyAsText() shouldContain "released" + } + } + + @Test + fun `a batch that repeats a request id is refused`() = testApplication { + application { mcpStreamableHttp { testServer() } } + val sessionId = client.initialize() + + val response = client.post("/mcp") { + localhostStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody("""[{"jsonrpc":"2.0","id":9,"method":"ping"},{"jsonrpc":"2.0","id":9,"method":"ping"}]""") + } + + response.status shouldBe HttpStatusCode.BadRequest + response.decodeError().error.code shouldBe RPCError.ErrorCode.INVALID_REQUEST + } + + @Test + fun `session limits must be positive`() { + assertFailsWith { + testApplication { + application { mcpStreamableHttp(maxSessions = 0) { testServer() } } + client.get("/mcp") + } + } + assertFailsWith { + testApplication { + application { mcpStreamableHttp(sessionIdleTimeout = Duration.ZERO) { testServer() } } + client.get("/mcp") + } + } + } + + private suspend fun HttpClient.initialize(): String { + val response = post("/mcp") { + localhostStreamableHeaders() + setBody(initializeRequestBody()) + } + response.status shouldBe HttpStatusCode.OK + return response.headers[MCP_SESSION_ID_HEADER].shouldNotBeNull() + } + + private suspend fun HttpClient.sendInitialized(sessionId: String) { + post("/mcp") { + localhostStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody("""{"jsonrpc":"2.0","method":"notifications/initialized"}""") + }.status shouldBe HttpStatusCode.Accepted + } + + private suspend fun HttpClient.ping(sessionId: String): HttpResponse = post("/mcp") { + localhostStreamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody("""{"jsonrpc":"2.0","id":2,"method":"ping"}""") + } + + /** [streamableHeaders] plus the `Host` that the endpoint's DNS rebinding protection admits. */ + private fun HttpRequestBuilder.localhostStreamableHeaders() { + header(HttpHeaders.Host, "localhost") + streamableHeaders() + } + + private fun initializeRequestBody(): String = + McpJson.encodeToString(JSONRPCMessage.serializer(), initializeRequest()) + + private suspend fun HttpResponse.decodeError(): JSONRPCError = McpJson.decodeFromString(bodyAsText()) +} diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionManagerTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionManagerTest.kt new file mode 100644 index 000000000..2144c2208 --- /dev/null +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpSessionManagerTest.kt @@ -0,0 +1,299 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.kotest.matchers.shouldBe +import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport +import io.modelcontextprotocol.kotlin.sdk.shared.Transport +import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.joinAll +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.test.advanceTimeBy +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.yield +import java.util.concurrent.atomic.AtomicBoolean +import kotlin.test.Test +import kotlin.time.Duration +import kotlin.time.Duration.Companion.days +import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.minutes +import kotlin.time.Duration.Companion.nanoseconds +import kotlin.time.Duration.Companion.seconds +import kotlin.time.TestTimeSource + +class StreamableHttpSessionManagerTest { + + private val timeSource = TestTimeSource() + private val idleTimeout = 10.minutes + + private fun sessionManager(maxSessions: Int = 10) = StreamableHttpSessionManager( + idleTimeout = idleTimeout, + maxSessions = maxSessions, + timeSource = timeSource, + ) + + private fun newTransport() = StreamableHttpServerTransport(StreamableHttpServerTransport.Configuration()) + + /** Opens a session the way an initialize request does, then completes that request. */ + private fun StreamableHttpSessionManager.openIdleSession( + sessionId: String, + transport: Transport = newTransport(), + ): Transport { + tryReserve() shouldBe true + register(sessionId, transport) + release(sessionId, transport) + return transport + } + + @Test + fun `idle session is closed once the idle timeout elapses`() = runTest { + val sessions = sessionManager() + val transport = sessions.openIdleSession("session") + var transportClosed = false + transport.onClose { transportClosed = true } + + timeSource += idleTimeout - 1.milliseconds + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe true + transportClosed shouldBe false + + timeSource += 1.milliseconds + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe false + transportClosed shouldBe true + } + + @Test + fun `session does not expire while its initializing request is in flight`() = runTest { + val sessions = sessionManager() + sessions.tryReserve() shouldBe true + val transport = newTransport() + sessions.register("session", transport) + + timeSource += idleTimeout * 2 + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe true + + sessions.release("session", transport) + timeSource += idleTimeout + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe false + } + + @Test + fun `session does not expire while a request is in flight`() = runTest { + val sessions = sessionManager() + sessions.openIdleSession("session") + + sessions.withSession("session") { + // A long-running request. + timeSource += idleTimeout * 2 + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe true + } shouldBe true + + timeSource += idleTimeout + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe false + } + + @Test + fun `a request restarts the idle timer`() = runTest { + val sessions = sessionManager() + sessions.openIdleSession("session") + + timeSource += idleTimeout - 1.milliseconds + sessions.withSession("session") {} shouldBe true + + timeSource += idleTimeout - 1.milliseconds + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe true + + timeSource += 1.milliseconds + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe false + } + + @Test + fun `touch restarts the idle timer without keeping the session open`() = runTest { + val sessions = sessionManager() + val transport = sessions.openIdleSession("session") + + timeSource += idleTimeout - 1.milliseconds + // A GET stream opens and stays open. + sessions.touch("session") shouldBe transport + + timeSource += idleTimeout - 1.milliseconds + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe true + + timeSource += 1.milliseconds + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe false + } + + @Test + fun `an unknown or removed session cannot be reached`() = runTest { + val sessions = sessionManager() + var ran = false + sessions.touch("unknown") shouldBe null + sessions.withSession("unknown") { ran = true } shouldBe false + + val transport = sessions.openIdleSession("session") + sessions.remove("session", transport) + sessions.touch("session") shouldBe null + sessions.withSession("session") { ran = true } shouldBe false + ran shouldBe false + } + + @Test + fun `new sessions are refused while maxSessions are open`() { + val sessions = sessionManager(maxSessions = 2) + sessions.tryReserve() shouldBe true + sessions.tryReserve() shouldBe true + sessions.tryReserve() shouldBe false + + sessions.cancelReservation() + sessions.tryReserve() shouldBe true + } + + @Test + fun `removing a session frees its slot once`() { + val sessions = sessionManager(maxSessions = 2) + val transport = sessions.openIdleSession("first") + sessions.openIdleSession("second") + sessions.tryReserve() shouldBe false + + sessions.remove("first", transport) + sessions.remove("first", transport) + ("first" in sessions) shouldBe false + sessions.tryReserve() shouldBe true + sessions.tryReserve() shouldBe false + } + + @Test + fun `expired session frees its slot`() = runTest { + val sessions = sessionManager(maxSessions = 1) + sessions.openIdleSession("session") + sessions.tryReserve() shouldBe false + + timeSource += idleTimeout + sessions.closeExpiredSessions() + sessions.tryReserve() shouldBe true + } + + @Test + fun `remove ignores a transport that does not serve the session`() { + val sessions = sessionManager() + sessions.openIdleSession("session") + + sessions.remove("session", newTransport()) + ("session" in sessions) shouldBe true + } + + @Test + fun `a close that hangs or fails does not hold up expiry`() = runTest { + val sessions = sessionManager() + val hanging = HangingTransport() + sessions.openIdleSession("hanging", hanging) + sessions.openIdleSession("failing", FailingTransport()) + var closed = false + sessions.openIdleSession("session").onClose { closed = true } + timeSource += idleTimeout + + val expiry = launch { sessions.closeExpiredSessionsPeriodically() } + try { + advanceTimeBy(sessions.expiryCheckInterval + 1.milliseconds) + hanging.closing.isCompleted shouldBe true + closed shouldBe true + + // Expiry carries on while that close still hangs. + var laterClosed = false + sessions.openIdleSession("later").onClose { laterClosed = true } + timeSource += idleTimeout + advanceTimeBy(sessions.expiryCheckInterval) + laterClosed shouldBe true + } finally { + expiry.cancel() + } + } + + @Test + fun `infinite idle timeout never expires sessions`() = runTest { + val sessions = StreamableHttpSessionManager(idleTimeout = Duration.INFINITE, timeSource = timeSource) + sessions.expires shouldBe false + sessions.openIdleSession("session") + + timeSource += 365.days + sessions.closeExpiredSessions() + ("session" in sessions) shouldBe true + } + + @Test + fun `expiry is checked at least every five seconds`() { + StreamableHttpSessionManager(idleTimeout = 30.minutes).expiryCheckInterval shouldBe 5.seconds + StreamableHttpSessionManager(idleTimeout = 1.seconds).expiryCheckInterval shouldBe 1.seconds + StreamableHttpSessionManager(idleTimeout = 1.nanoseconds).expiryCheckInterval shouldBe + 10.milliseconds + } + + @Test + fun `expiry never closes a session while a request runs on it`(): Unit = runBlocking(Dispatchers.Default) { + // Every idle moment is past the timeout, so expiry races the start of every request. + val sessions = StreamableHttpSessionManager(idleTimeout = 1.nanoseconds, maxSessions = 1) + repeat(200) { round -> + val sessionId = "session-$round" + val transport = newTransport() + val transportClosed = AtomicBoolean(false) + transport.onClose { transportClosed.set(true) } + sessions.openIdleSession(sessionId, transport) + + val requestRunning = CompletableDeferred() + val requests = List(4) { + launch { + repeat(50) { + sessions.withSession(sessionId) { + requestRunning.complete(Unit) + transportClosed.get() shouldBe false + yield() + transportClosed.get() shouldBe false + } + } + } + } + // Start expiring only once a request runs, so that expiry races requests instead of beating them all. + requestRunning.await() + while (sessionId in sessions) { + sessions.closeExpiredSessions() + yield() + } + requests.joinAll() + transportClosed.get() shouldBe true + } + } + + /** A transport whose close never finishes, like one stuck on a client that stopped reading. */ + private class HangingTransport : AbstractTransport() { + val closing = CompletableDeferred() + + override suspend fun start() = Unit + + override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) = Unit + + override suspend fun close() { + closing.complete(Unit) + awaitCancellation() + } + } + + private class FailingTransport : AbstractTransport() { + override suspend fun start() = Unit + + override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) = Unit + + override suspend fun close(): Unit = error("close failed") + } +} diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpStreamReleaseTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpStreamReleaseTest.kt new file mode 100644 index 000000000..99da758d9 --- /dev/null +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpStreamReleaseTest.kt @@ -0,0 +1,247 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.kotest.matchers.nulls.shouldNotBeNull +import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldContain +import io.ktor.client.HttpClient +import io.ktor.client.request.header +import io.ktor.client.request.post +import io.ktor.client.request.prepareGet +import io.ktor.client.request.setBody +import io.ktor.client.statement.HttpStatement +import io.ktor.client.statement.bodyAsChannel +import io.ktor.client.statement.bodyAsText +import io.ktor.http.ContentType +import io.ktor.http.HttpHeaders +import io.ktor.http.HttpStatusCode +import io.ktor.server.application.install +import io.ktor.server.routing.post +import io.ktor.server.routing.routing +import io.ktor.server.sse.SSE +import io.ktor.server.sse.ServerSSESession +import io.ktor.server.sse.sse +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.modelcontextprotocol.kotlin.sdk.types.EmptyResult +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 kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.delay +import kotlinx.coroutines.isActive +import kotlinx.coroutines.launch +import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock +import kotlinx.coroutines.withTimeout +import kotlinx.coroutines.withTimeoutOrNull +import kotlin.test.Test +import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.seconds + +/** + * Verifies that a standalone GET handler returns once the transport stops using its stream, and that a closed + * transport refuses requests instead of leaving them waiting. An engine may never cancel a handler that is only + * waiting, even after its client went away, so a handler that waited for cancellation stayed alive for good. + */ +class StreamableHttpStreamReleaseTest { + + private val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration(enableJsonResponse = true), + ).apply { + onMessage { message -> + if (message is JSONRPCRequest) send(JSONRPCResponse(message.id, EmptyResult())) + } + } + + /** Receives the [STREAM_NAME_HEADER] of each GET whose handler finishes, by returning or by throwing. */ + private val finishedStreams = Channel(Channel.UNLIMITED) + + private fun ApplicationTestBuilder.serveTransport(wrap: (ServerSSESession) -> ServerSSESession = { it }) { + application { + installMcpContentNegotiation() + install(SSE) + routing { + post("/mcp") { transport.handleRequest(null, call) } + sse("/mcp") { + try { + transport.handleRequest(wrap(this), call) + } finally { + finishedStreams.trySend(call.request.headers[STREAM_NAME_HEADER].orEmpty()) + } + } + } + } + } + + @Test + fun `closing the transport ends an open GET stream`() = testApplication { + serveTransport() + val sessionId = client.initialize() + + client.openStream(sessionId).execute { stream -> + stream.bodyAsChannel().readLine() + + transport.close() + + withTimeout(5.seconds) { finishedStreams.receive() } + } + } + + @Test + fun `a newer GET stream ends the one it replaces`() = testApplication { + serveTransport() + val sessionId = client.initialize() + + client.openStream(sessionId, name = "first").execute { first -> + awaitRoutedTo(first.bodyAsChannel()) + + client.openStream(sessionId, name = "second").execute { second -> + second.bodyAsChannel().readLine() + + withTimeout(5.seconds) { finishedStreams.receive() } shouldBe "first" + finishedStreams.tryReceive().isSuccess shouldBe false + } + } + } + + @Test + fun `closing does not wait on a stream whose client stopped reading`() = testApplication { + val stalledSessions = Channel(Channel.UNLIMITED) + serveTransport { session -> StalledSseSession(session).also { stalledSessions.trySend(it) } } + val sessionId = client.initialize() + + client.openStream(sessionId).execute { stream -> + stream.bodyAsChannel().readLine() + val session = stalledSessions.receive() + session.stall() + + coroutineScope { + // Server-initiated messages until one gets stuck on the stream, holding the lock that closing + // the SSE session takes, as a send to a client that stopped reading does. + val sender = launch { + while (isActive) { + transport.send(JSONRPCNotification(method = "notifications/message")) + delay(10.milliseconds) + } + } + try { + withTimeout(5.seconds) { session.sendStuck.await() } + + val closing = launch { transport.close() } + // A close that waited for the stuck send would never finish, so join with a deadline. + withTimeoutOrNull(5.seconds) { closing.join() }.shouldNotBeNull() + withTimeout(5.seconds) { finishedStreams.receive() } + } finally { + sender.cancel() + } + } + } + } + + @Test + fun `POST on a closed transport is rejected instead of left waiting`() = testApplication { + serveTransport() + val sessionId = client.initialize() + transport.close() + + val response = withTimeout(5.seconds) { + client.post("/mcp") { + streamableHeaders() + header(MCP_SESSION_ID_HEADER, sessionId) + setBody("""{"jsonrpc":"2.0","id":2,"method":"ping"}""") + } + } + response.status shouldBe HttpStatusCode.NotFound + response.bodyAsText() shouldContain "Session not found" + } + + private suspend fun HttpClient.initialize(): String { + val response = post("/mcp") { + streamableHeaders() + setBody(McpJson.encodeToString(JSONRPCMessage.serializer(), initializeRequest())) + } + response.status shouldBe HttpStatusCode.OK + return response.headers[MCP_SESSION_ID_HEADER].shouldNotBeNull() + } + + /** Returns once the transport routes server-initiated messages to the stream read through [channel]. */ + private suspend fun awaitRoutedTo(channel: ByteReadChannel) = coroutineScope { + // Messages sent before the stream registers are dropped, so keep probing until one arrives. + val prober = launch { + while (isActive) { + transport.send(JSONRPCNotification(method = "notifications/probe")) + delay(10.milliseconds) + } + } + try { + withTimeout(5.seconds) { + do { + val line = channel.readLine() + } while (line?.contains("notifications/probe") != true) + } + } finally { + prober.cancel() + } + } + + private suspend fun HttpClient.openStream(sessionId: String, name: String = ""): HttpStatement = + prepareGet("/mcp") { + header(HttpHeaders.Accept, ContentType.Text.EventStream.toString()) + header(MCP_SESSION_ID_HEADER, sessionId) + header(STREAM_NAME_HEADER, name) + } + + /** + * An SSE session whose client stops reading once [stall] is called: a send then never completes and keeps holding + * the lock that [close] also takes, like Ktor's own session does while its buffer is full. + */ + private class StalledSseSession(private val delegate: ServerSSESession) : ServerSSESession by delegate { + private val lock = Mutex() + + @Volatile + private var stalled = false + + val sendStuck = CompletableDeferred() + + fun stall() { + stalled = true + } + + override suspend fun send(event: ServerSentEvent) { + lock.withLock { + if (stalled) { + sendStuck.complete(Unit) + awaitCancellation() + } + delegate.send(event) + } + } + + // Delegation would route this default method to the delegate, bypassing the stall. + override suspend fun send( + data: String?, + event: String?, + id: String?, + retry: Long?, + comments: String?, + ) { + send(ServerSentEvent(data, event, id, retry, comments)) + } + + override suspend fun close() { + lock.withLock { delegate.close() } + } + } + + private companion object { + const val STREAM_NAME_HEADER = "X-Test-Stream" + } +}