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 dcba84873..f4fc4aca1 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 @@ -59,7 +59,16 @@ private const val MIN_PRIMING_EVENT_PROTOCOL_VERSION = "2025-11-25" * 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 data class SessionContext(val session: ServerSSESession?, val call: ApplicationCall) { + /** + * Cancels the coroutine serving this stream's HTTP call, releasing the underlying + * connection. Required for streams whose handler suspends for the lifetime of the + * stream, which closing the [ServerSSESession] alone does not interrupt. + */ + fun cancelCall() { + call.coroutineContext.job.cancel() + } +} /** * Server transport for Streamable HTTP: this implements the MCP Streamable HTTP transport specification. @@ -383,6 +392,11 @@ public class StreamableHttpServerTransport(private val configuration: Configurat } catch (_: Exception) { } } + // Closing the SSE session only closes the response body. The standalone GET + // handler is parked on awaitCancellation(), so its call has to be cancelled + // explicitly; otherwise that coroutine — and the connection behind it — is + // never released. + streamsMapping[STANDALONE_SSE_STREAM_ID]?.cancelCall() streamsMapping.clear() requestToStreamMapping.clear() requestToResponseMapping.clear() @@ -613,9 +627,10 @@ public class StreamableHttpServerTransport(private val configuration: Configurat 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). + // Close the previous SSE session, then cancel its call. Closing the session + // only closes the response body, while the previous handler is parked on + // awaitCancellation(); cancelling makes it hit its identity-guarded finally — + // that finally won't double-remove, since we replace the mapping below. try { existingContext.session?.close() } catch (e: CancellationException) { @@ -623,6 +638,7 @@ public class StreamableHttpServerTransport(private val configuration: Configurat } catch (_: Exception) { // Ignore — the old stream may already be closed. } + existingContext.cancelCall() // Evict the stale mapping — the old session is closed either way. streamsMapping.remove(STANDALONE_SSE_STREAM_ID) } diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StatefulStreamableHttpGetStreamLifecycleTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StatefulStreamableHttpGetStreamLifecycleTest.kt new file mode 100644 index 000000000..f91516599 --- /dev/null +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StatefulStreamableHttpGetStreamLifecycleTest.kt @@ -0,0 +1,169 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.kotest.matchers.shouldBe +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.http.ContentType +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.testApplication +import io.ktor.sse.ServerSentEvent +import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest +import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION +import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.toJSON +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.cancel +import kotlinx.coroutines.launch +import kotlinx.coroutines.withTimeoutOrNull +import kotlin.coroutines.CoroutineContext +import kotlin.test.Test +import kotlin.test.assertNotNull +import kotlin.time.Duration.Companion.seconds +import io.ktor.server.plugins.contentnegotiation.ContentNegotiation as ServerContentNegotiation + +/** + * The standalone GET stream keeps its handler suspended on `awaitCancellation()` for the + * lifetime of the stream. Closing the [ServerSSESession] only closes the response body, so + * tearing the session down has to cancel that call as well — otherwise the coroutine and the + * connection behind it are never released + * (https://github.com/modelcontextprotocol/kotlin-sdk/issues/922). + */ +class StatefulStreamableHttpGetStreamLifecycleTest { + + @Test + fun `closing the transport releases the standalone GET stream handler`() = testApplication { + val transport = + StreamableHttpServerTransport(StreamableHttpServerTransport.Configuration(enableJsonResponse = true)) + val handlerReleased = CompletableDeferred() + + val sessionId = startSessionWithGetStream(transport, handlerReleased) + + transport.close() + + awaitRelease(handlerReleased) shouldBe true + assertNotNull(sessionId) + } + + @Test + fun `a replacement GET stream releases the previous stream handler`() = testApplication { + val transport = + StreamableHttpServerTransport(StreamableHttpServerTransport.Configuration(enableJsonResponse = true)) + val firstReleased = CompletableDeferred() + + val sessionId = assertNotNull(startSessionWithGetStream(transport, firstReleased)) + + // A client reconnecting its GET stream takes over the standalone slot; the handler + // serving the previous stream must be released rather than left suspended. + val scope = CoroutineScope(Dispatchers.Default) + scope.launch { openGetStream(sessionId) } + + awaitRelease(firstReleased) shouldBe true + scope.cancel() + } + + /** + * Initializes a session against [transport] and opens the standalone GET stream, completing + * [handlerReleased] when the GET handler returns. Fails if the handler returns before the + * stream is torn down, which would make the assertions vacuous. + */ + private suspend fun io.ktor.server.testing.ApplicationTestBuilder.startSessionWithGetStream( + transport: StreamableHttpServerTransport, + handlerReleased: CompletableDeferred, + ): String? { + val streamOpened = CompletableDeferred() + application { + install(ServerContentNegotiation) { json(McpJson) } + routing { + post(PATH) { transport.handlePostRequest(null, call) } + get(PATH) { + try { + transport.handleGetRequest(RecordingSseSession(call, streamOpened), call) + } finally { + handlerReleased.complete(Unit) + } + } + } + } + Server( + Implementation("test-server", "1.0.0"), + ServerOptions(capabilities = ServerCapabilities()), + ).createSession(transport) + + val initResponse = client.post(PATH) { + header(HttpHeaders.Host, "localhost") + header( + HttpHeaders.Accept, + listOf(ContentType.Application.Json, ContentType.Text.EventStream).joinToString(", "), + ) + contentType(ContentType.Application.Json) + setBody(McpJson.encodeToString(JSONRPCMessage.serializer(), initializePayload())) + } + initResponse.status shouldBe HttpStatusCode.OK + val sessionId = initResponse.headers[MCP_SESSION_ID_HEADER] + + val scope = CoroutineScope(Dispatchers.Default) + scope.launch { openGetStream(assertNotNull(sessionId)) } + + withTimeoutOrNull(5.seconds) { streamOpened.await() } + streamOpened.isCompleted shouldBe true + // The handler must still be suspended at this point, otherwise the test proves nothing. + handlerReleased.isCompleted shouldBe false + return sessionId + } + + private suspend fun io.ktor.server.testing.ApplicationTestBuilder.openGetStream(sessionId: String) { + client.get(PATH) { + header(HttpHeaders.Host, "localhost") + header(HttpHeaders.Accept, ContentType.Text.EventStream.toString()) + header(MCP_SESSION_ID_HEADER, sessionId) + header("mcp-protocol-version", LATEST_PROTOCOL_VERSION) + } + } + + private suspend fun awaitRelease(released: CompletableDeferred): Boolean = withTimeoutOrNull(5.seconds) { + released.await() + true + } ?: false + + private fun initializePayload() = InitializeRequest( + InitializeRequestParams( + protocolVersion = LATEST_PROTOCOL_VERSION, + capabilities = ClientCapabilities(), + clientInfo = Implementation(name = "test-client", version = "1.0.0"), + ), + ).toJSON() + + private companion object { + const val PATH = "/mcp" + } +} + +/** Signals [opened] once the transport writes to the stream, i.e. the handler is about to suspend. */ +private class RecordingSseSession(override val call: ApplicationCall, private val opened: CompletableDeferred) : + ServerSSESession { + override val coroutineContext: CoroutineContext = call.coroutineContext + + override suspend fun send(event: ServerSentEvent) { + opened.complete(Unit) + } + + override suspend fun close() {} +}