From a06d699d6c9dfdf8321e5cbb3f5aebcb0016b6f6 Mon Sep 17 00:00:00 2001 From: UnscientificJsZhai <73090853+UnscientificJsZhai@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:37:37 +0800 Subject: [PATCH] feat: SSE heartbeat (#761) ## Summary Adds optional SSE heartbeat configuration for streamable HTTP MCP server connections. This lets `Application.mcpStreamableHttp` callers pass Ktor's `Heartbeat` configuration for SSE streams, including heartbeat period and event payload. Heartbeats remain disabled by default, preserving the existing stream behavior unless the new option is explicitly provided. ## Motivation and Context Some MCP clients need heartbeat messages on SSE connections to keep the connection alive. Without those heartbeats, the client may proactively disconnect from the MCP server. MCP server developers need a way to configure an SSE heartbeat mechanism so they can avoid unexpected client disconnects. In my case, my MCP server was consistently disconnected by Gemini CLI. I worked around the issue with some fallback approaches, but I believe the best implementation is to add heartbeat support by using Ktor's built-in API. ## Implementation Notes - Added a function overload of `Application.mcpStreamableHttp` to introduce heartbeat parameter `sseHeartbeatConfig: (Heartbeat.() -> Unit)?`. - Applied the configuration inside the Ktor `sse { ... }` block before handling the existing streamable HTTP transport request. - Updated the server API dump for the new public API surface. ## How Has This Been Tested? Added `StreamableHttpHeartbeatIntegrationTest` covering: - Configured GET SSE streams apply the provided heartbeat configuration. - GET SSE streams do not send default heartbeats when `sseHeartbeatConfig` is omitted. - Configured GET SSE streams emit heartbeat events repeatedly. - Legacy MCP over SSE APIs do not expose the new heartbeat configuration. Verified locally with: ```bash ./gradlew :integration-test:jvmTest --tests "io.modelcontextprotocol.kotlin.sdk.integration.streamablehttp.StreamableHttpHeartbeatIntegrationTest" ``` Result: `BUILD SUCCESSFUL`. ## Breaking Changes None. This change adds a new `Application.mcpStreamableHttp` overload that accepts `sseHeartbeatConfig: (Heartbeat.() -> Unit)?` while keeping the existing overload unchanged. Existing callers continue to compile and keep the previous behavior: SSE heartbeats are not sent unless the new overload is used with an explicit heartbeat configuration. ## Types of changes - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update ## Checklist - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [x] New and existing tests pass locally - [x] I have added appropriate error handling - [x] I have added or updated documentation as needed ## Additional context Heartbeat behavior is delegated to Ktor's SSE heartbeat support, so applications can use the same configuration semantics they would use in native Ktor SSE routes. Related issue: [SSE needs a heartbeat #344 ](https://github.com/modelcontextprotocol/kotlin-sdk/issues/344) --------- Co-authored-by: Pavel Gorgulov --- .../AbstractStreamableHttpIntegrationTest.kt | 8 +- .../StreamableHttpHeartbeatIntegrationTest.kt | 174 ++++++++++++++++++ kotlin-sdk-server/api/kotlin-sdk-server.api | 4 +- .../kotlin/sdk/server/KtorServer.kt | 7 + 4 files changed, 189 insertions(+), 4 deletions(-) create mode 100644 integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpHeartbeatIntegrationTest.kt diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/AbstractStreamableHttpIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/AbstractStreamableHttpIntegrationTest.kt index 9bc0c17..535393e 100644 --- a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/AbstractStreamableHttpIntegrationTest.kt +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/AbstractStreamableHttpIntegrationTest.kt @@ -3,6 +3,7 @@ package io.modelcontextprotocol.kotlin.sdk.integration.streamablehttp import io.ktor.server.cio.CIOApplicationEngine import io.ktor.server.engine.EmbeddedServer import io.ktor.server.engine.embeddedServer +import io.ktor.server.sse.Heartbeat import io.modelcontextprotocol.kotlin.sdk.server.Server import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions import io.modelcontextprotocol.kotlin.sdk.server.mcpStreamableHttp @@ -17,7 +18,10 @@ import io.ktor.server.cio.CIO as ServerCIO open class AbstractStreamableHttpIntegrationTest { - suspend fun initTestServer(name: String? = null): StreamableHttpTestServer { + suspend fun initTestServer( + name: String? = null, + sseHeartbeatConfig: (Heartbeat.() -> Unit)? = null, + ): StreamableHttpTestServer { val mcpServer = Server( Implementation(name = name ?: DEFAULT_SERVER_NAME, version = VERSION), ServerOptions( @@ -52,7 +56,7 @@ open class AbstractStreamableHttpIntegrationTest { host = URL, port = PORT, ) { - mcpStreamableHttp { mcpServer } + mcpStreamableHttp(sseHeartbeatConfig = sseHeartbeatConfig) { mcpServer } } return StreamableHttpTestServer( diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpHeartbeatIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpHeartbeatIntegrationTest.kt new file mode 100644 index 0000000..8019557 --- /dev/null +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpHeartbeatIntegrationTest.kt @@ -0,0 +1,174 @@ +package io.modelcontextprotocol.kotlin.sdk.integration.streamablehttp + +import io.kotest.matchers.shouldBe +import io.ktor.client.HttpClient +import io.ktor.client.plugins.sse.SSE +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.bodyAsChannel +import io.ktor.http.ContentType +import io.ktor.http.HttpHeaders +import io.ktor.http.HttpStatusCode +import io.ktor.http.contentType +import io.ktor.sse.ServerSentEvent +import io.ktor.utils.io.ByteReadChannel +import io.ktor.utils.io.readUTF8Line +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.JSONRPCRequest +import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION +import io.modelcontextprotocol.kotlin.sdk.types.toJSON +import io.modelcontextprotocol.kotlin.test.utils.actualPort +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeoutOrNull +import kotlinx.serialization.json.Json +import kotlin.test.Test +import kotlin.test.assertNotNull +import kotlin.time.Duration.Companion.milliseconds +import io.ktor.client.engine.cio.CIO as ClientCIO + +private const val SESSION_ID_HEADER = "mcp-session-id" +private const val PROTOCOL_VERSION_HEADER = "mcp-protocol-version" +private const val DEFAULT_HEARTBEAT_LINE = ": heartbeat" +private const val CUSTOM_HEARTBEAT_LINE = ": mcp-heartbeat" + +class StreamableHttpHeartbeatIntegrationTest : AbstractStreamableHttpIntegrationTest() { + + @Test + fun `GET SSE stream emits configured heartbeat`(): Unit = runBlocking(Dispatchers.IO) { + var server: StreamableHttpTestServer? = null + var httpClient: HttpClient? = null + + try { + server = initTestServer("heartbeat-test") { + period = 50.milliseconds + event = ServerSentEvent(comments = "mcp-heartbeat") + } + val mcpUrl = "http://$URL:${server.ktorServer.actualPort()}/mcp" + httpClient = HttpClient(ClientCIO) { install(SSE) } + val sessionId = initializeSession(httpClient, mcpUrl) + + httpClient.prepareGet(mcpUrl) { + addSseHeaders(sessionId) + }.execute { response -> + response.status shouldBe HttpStatusCode.OK + response.headers[SESSION_ID_HEADER] shouldBe sessionId + + response.bodyAsChannel().readLineMatching(CUSTOM_HEARTBEAT_LINE) shouldBe CUSTOM_HEARTBEAT_LINE + } + } finally { + httpClient?.close() + server?.ktorServer?.stopSuspend(1000, 2000) + } + } + + @Test + fun `GET SSE stream does not emit heartbeat by default`(): Unit = runBlocking(Dispatchers.IO) { + var server: StreamableHttpTestServer? = null + var httpClient: HttpClient? = null + + try { + server = initTestServer("no-heartbeat-test") + val mcpUrl = "http://$URL:${server.ktorServer.actualPort()}/mcp" + httpClient = HttpClient(ClientCIO) { install(SSE) } + val sessionId = initializeSession(httpClient, mcpUrl) + + httpClient.prepareGet(mcpUrl) { + addSseHeaders(sessionId) + }.execute { response -> + response.status shouldBe HttpStatusCode.OK + response.headers[SESSION_ID_HEADER] shouldBe sessionId + + response.bodyAsChannel().readLineMatching(timeoutMillis = 150) { line -> + line.isHeartbeatSseLine() + } shouldBe null + } + } finally { + httpClient?.close() + server?.ktorServer?.stopSuspend(1000, 2000) + } + } + + @Test + fun `GET SSE stream emits configured heartbeat repeatedly`(): Unit = runBlocking(Dispatchers.IO) { + var server: StreamableHttpTestServer? = null + var httpClient: HttpClient? = null + + try { + server = initTestServer("repeating-heartbeat-test") { + period = 50.milliseconds + event = ServerSentEvent(comments = "mcp-heartbeat") + } + val mcpUrl = "http://$URL:${server.ktorServer.actualPort()}/mcp" + httpClient = HttpClient(ClientCIO) { install(SSE) } + val sessionId = initializeSession(httpClient, mcpUrl) + + httpClient.prepareGet(mcpUrl) { + addSseHeaders(sessionId) + }.execute { response -> + response.status shouldBe HttpStatusCode.OK + response.headers[SESSION_ID_HEADER] shouldBe sessionId + + val channel = response.bodyAsChannel() + channel.readLineMatching(CUSTOM_HEARTBEAT_LINE) shouldBe CUSTOM_HEARTBEAT_LINE + channel.readLineMatching(CUSTOM_HEARTBEAT_LINE) shouldBe CUSTOM_HEARTBEAT_LINE + } + } finally { + httpClient?.close() + server?.ktorServer?.stopSuspend(1000, 2000) + } + } + + private suspend fun initializeSession(client: HttpClient, mcpUrl: String): String { + val response = client.post(mcpUrl) { + contentType(ContentType.Application.Json) + header( + HttpHeaders.Accept, + "${ContentType.Application.Json}, ${ContentType.Text.EventStream}", + ) + setBody(Json.encodeToString(buildInitPayload())) + } + + response.status shouldBe HttpStatusCode.OK + return assertNotNull(response.headers[SESSION_ID_HEADER]) + } + + private fun buildInitPayload(): JSONRPCRequest = InitializeRequest( + InitializeRequestParams( + protocolVersion = LATEST_PROTOCOL_VERSION, + capabilities = ClientCapabilities(), + clientInfo = Implementation(name = "heartbeat-test-client", version = "1.0.0"), + ), + ).toJSON() + + private fun io.ktor.client.request.HttpRequestBuilder.addSseHeaders(sessionId: String) { + header(HttpHeaders.Accept, ContentType.Text.EventStream.toString()) + header(SESSION_ID_HEADER, sessionId) + header(PROTOCOL_VERSION_HEADER, LATEST_PROTOCOL_VERSION) + } + + private suspend fun ByteReadChannel.readLineMatching(expectedLine: String, timeoutMillis: Long = 2_000): String? = + readLineMatching(timeoutMillis) { line -> line == expectedLine } + + private suspend fun ByteReadChannel.readLineMatching( + timeoutMillis: Long = 2_000, + matches: (String) -> Boolean, + ): String? = withTimeoutOrNull(timeoutMillis.milliseconds) { + var line = readUTF8Line() + while (line != null) { + if (matches(line)) return@withTimeoutOrNull line + line = readUTF8Line() + } + null + } + + private fun String.isHeartbeatSseLine(): Boolean = this == DEFAULT_HEARTBEAT_LINE || + this == CUSTOM_HEARTBEAT_LINE || + (startsWith(":") && contains("heartbeat", ignoreCase = true)) || + (startsWith("event:") && contains("heartbeat", ignoreCase = true)) +} diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index 51ac279..1ec2efd 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -59,8 +59,8 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/KtorServerKt { public static synthetic fun mcp$default (Lio/ktor/server/routing/Route;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static final fun mcpStatelessStreamableHttp (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;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 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;)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;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 final class io/modelcontextprotocol/kotlin/sdk/server/RegisteredPrompt : io/modelcontextprotocol/kotlin/sdk/server/Feature { 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 73ef258..5fc3d22 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 @@ -20,8 +20,10 @@ import io.ktor.server.routing.get import io.ktor.server.routing.post import io.ktor.server.routing.route import io.ktor.server.routing.routing +import io.ktor.server.sse.Heartbeat import io.ktor.server.sse.SSE import io.ktor.server.sse.ServerSSESession +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 @@ -162,6 +164,7 @@ private fun Application.mcpStreamableHttp( allowedHosts: List?, allowedOrigins: List?, configuration: StreamableHttpServerTransport.Configuration, + sseHeartbeatConfig: (Heartbeat.() -> Unit)?, block: RoutingContext.() -> Server, ) { installMcpContentNegotiation() @@ -185,6 +188,7 @@ private fun Application.mcpStreamableHttp( sse { val transport = existingStreamableTransport(call, transportManager) ?: return@sse + sseHeartbeatConfig?.let { config -> heartbeat(config) } transport.handleRequest(this, call) } @@ -227,6 +231,7 @@ private fun Application.mcpStreamableHttp( * With custom `allowedHosts`, `null` skips origin validation. * @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 block factory block with access to the [RoutingContext] (for reading request headers) * that creates and returns the [Server] to handle the connection. */ @@ -237,6 +242,7 @@ public fun Application.mcpStreamableHttp( allowedHosts: List? = null, allowedOrigins: List? = null, eventStore: EventStore? = null, + sseHeartbeatConfig: (Heartbeat.() -> Unit)? = null, block: RoutingContext.() -> Server, ) { mcpStreamableHttp( @@ -248,6 +254,7 @@ public fun Application.mcpStreamableHttp( eventStore = eventStore, enableJsonResponse = true, ), + sseHeartbeatConfig = sseHeartbeatConfig, block = block, ) }