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..465f70fe3 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 @@ -34,6 +34,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.NonCancellable import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.currentCoroutineContext import kotlinx.coroutines.job import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock @@ -375,9 +376,13 @@ public class StreamableHttpServerTransport(private val configuration: Configurat } override suspend fun close() { + val currentJob = currentCoroutineContext().job withContext(NonCancellable) { streamMutex.withLock { streamsMapping.values.forEach { + if (it.call.coroutineContext.job !== currentJob) { + it.call.coroutineContext.job.cancel() + } try { it.session?.close() } catch (_: Exception) { @@ -611,12 +616,16 @@ public class StreamableHttpServerTransport(private val configuration: Configurat // SSE headers (Content-Type, Cache-Control, Connection) are already set by the framework's SSE handler flushSse(sseSession) val newContext = SessionContext(sseSession, call) + val currentJob = currentCoroutineContext().job 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 { + if (existingContext.call.coroutineContext.job !== currentJob) { + existingContext.call.coroutineContext.job.cancel() + } existingContext.session?.close() } catch (e: CancellationException) { throw e @@ -660,9 +669,13 @@ public class StreamableHttpServerTransport(private val configuration: Configurat if (configuration.enableJsonResponse) return val streamId = requestToStreamMapping[requestId] ?: return val sessionContext = streamsMapping[streamId] ?: return + val currentJob = currentCoroutineContext().job withContext(NonCancellable) { try { + if (sessionContext.call.coroutineContext.job !== currentJob) { + sessionContext.call.coroutineContext.job.cancel() + } sessionContext.session?.close() } catch (e: Exception) { _onError(e) 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 43ff6051e..74a2a8e90 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 @@ -16,12 +16,14 @@ import io.ktor.client.request.prepareGet import io.ktor.client.request.setBody import io.ktor.client.statement.bodyAsChannel import io.ktor.http.ContentType +import io.ktor.http.Headers 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.request.ApplicationRequest import io.ktor.server.routing.get import io.ktor.server.routing.post import io.ktor.server.routing.routing @@ -32,6 +34,8 @@ import io.ktor.sse.ServerSentEvent import io.ktor.utils.io.ByteReadChannel import io.ktor.utils.io.readLine import io.ktor.utils.io.readUTF8Line +import io.mockk.every +import io.mockk.mockk import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotification import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities @@ -64,6 +68,7 @@ import kotlinx.coroutines.cancel import kotlinx.coroutines.delay import kotlinx.coroutines.launch import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withContext import kotlinx.coroutines.withTimeout import kotlinx.serialization.builtins.ListSerializer import kotlinx.serialization.json.buildJsonObject @@ -1065,6 +1070,64 @@ class StreamableHttpServerTransportTest { } } + @Test + fun `closing transport cancels standalone GET SSE request`() = runTest { + val callJob = SupervisorJob() + val callContext = callJob + Dispatchers.Default + val call = mockk(relaxed = true) + val request = mockk(relaxed = true) + val requestHeaders = Headers.build { + append(HttpHeaders.Accept, ContentType.Text.EventStream.toString()) + append("mcp-protocol-version", LATEST_PROTOCOL_VERSION) + } + every { call.coroutineContext } returns callContext + every { call.request } returns request + every { request.headers } returns requestHeaders + + val mappingRegistered = CompletableDeferred() + val session = FakeServerSSESession(call, callContext) + val transport = StreamableHttpServerTransport( + StreamableHttpServerTransport.Configuration( + eventStore = object : EventStore { + override suspend fun storeEvent(streamId: String, message: JSONRPCMessage): String { + mappingRegistered.complete(Unit) + return "priming-event" + } + + override suspend fun replayEventsAfter( + lastEventId: String, + sender: suspend (eventId: String, message: JSONRPCMessage) -> Unit, + ): String = "standalone-stream" + + override suspend fun getStreamIdForEventId(eventId: String): String? = null + }, + ), + ) + transport.setSessionIdGenerator(null) + + val handler = CoroutineScope(callContext).launch { + transport.handleGetRequest(session, call) + } + + try { + withContext(Dispatchers.Default) { + withTimeout(5.seconds) { mappingRegistered.await() } + } + + transport.close() + assertTrue(callJob.isCancelled) + + withContext(Dispatchers.Default) { + withTimeout(5.seconds) { handler.join() } + } + assertFalse(handler.isActive) + } finally { + callJob.cancel() + handler.cancel() + transport.close() + } + } + @Test fun `GET SSE reconnect after previous stream disconnects should succeed`() = testApplication { val mcpPath = "/mcp"