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