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 <devcrocod@gmail.com>
This commit is contained in:
UnscientificJsZhai
GitHub
parent
cab669b108
commit
a06d699d6c
4 changed files with 189 additions and 4 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue