feat: SSE heartbeat

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(

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))
+}

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 {

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<String>?,
     allowedOrigins: List<String>?,
     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<String>? = null,
     allowedOrigins: List<String>? = 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,
     )
 }