diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index d8e4c7a5e..08fb3c889 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -52,17 +52,25 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/HostValidationKt { public final class io/modelcontextprotocol/kotlin/sdk/server/KtorServerKt { public static final fun mcp (Lio/ktor/server/application/Application;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;)V + public static final fun mcp (Lio/ktor/server/application/Application;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)V public static final fun mcp (Lio/ktor/server/routing/Route;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;)V + public static final fun mcp (Lio/ktor/server/routing/Route;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)V public static final fun mcp (Lio/ktor/server/routing/Route;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;)V + public static final fun mcp (Lio/ktor/server/routing/Route;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)V public static synthetic fun mcp$default (Lio/ktor/server/application/Application;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public static synthetic fun mcp$default (Lio/ktor/server/application/Application;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun mcp$default (Lio/ktor/server/routing/Route;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public static synthetic fun mcp$default (Lio/ktor/server/routing/Route;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V 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 synthetic fun mcp$default (Lio/ktor/server/routing/Route;ZLjava/util/List;Ljava/util/List;JLkotlin/jvm/functions/Function1;Lkotlin/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 final fun mcpStatelessStreamableHttp (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;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 synthetic fun mcpStatelessStreamableHttp$default (Lio/ktor/server/application/Application;Ljava/lang/String;ZLjava/util/List;Ljava/util/List;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 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;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 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;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 31274995e..6a3cadd77 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 @@ -27,6 +27,7 @@ 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.shared.Transport import io.modelcontextprotocol.kotlin.sdk.types.RPCError import kotlinx.coroutines.awaitCancellation @@ -62,9 +63,53 @@ public fun Route.mcp( allowedOrigins: List? = null, maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE, block: ServerSSESession.() -> Server, +) { + mcp( + path = path, + enableDnsRebindingProtection = enableDnsRebindingProtection, + allowedHosts = allowedHosts, + allowedOrigins = allowedOrigins, + maxRequestBodySize = maxRequestBodySize, + transportDecorator = { it }, + block = block, + ) +} + +/** + * Registers MCP over SSE at [path], decorating each transport before it is connected to its [Server]. + * + * The returned [Transport] must delegate lifecycle methods and callback registration to the supplied + * [SseServerTransport] so that Ktor request handling and the MCP session remain connected. + * + * @param path the URL path to register the SSE endpoint. + * @param enableDnsRebindingProtection whether to install [DnsRebindingProtection] on this route. + * @param allowedHosts hostnames allowed in the `Host` header. + * @param allowedOrigins origins allowed in the `Origin` header, compared by hostname only. + * @param maxRequestBodySize maximum allowed size, in bytes, of an incoming POST body. + * @param transportDecorator transforms each SDK-created transport before [Server.createSession] receives it. + * @param block factory block that creates the [Server] for the connection. + * @throws IllegalStateException if the [SSE] plugin is not installed. + */ +@KtorDsl +@Suppress("LongParameterList") +public fun Route.mcp( + path: String, + enableDnsRebindingProtection: Boolean = true, + allowedHosts: List? = null, + allowedOrigins: List? = null, + maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE, + transportDecorator: (SseServerTransport) -> Transport, + block: ServerSSESession.() -> Server, ) { route(path) { - mcp(enableDnsRebindingProtection, allowedHosts, allowedOrigins, maxRequestBodySize, block) + mcp( + enableDnsRebindingProtection = enableDnsRebindingProtection, + allowedHosts = allowedHosts, + allowedOrigins = allowedOrigins, + maxRequestBodySize = maxRequestBodySize, + transportDecorator = transportDecorator, + block = block, + ) } } @@ -96,6 +141,40 @@ public fun Route.mcp( allowedOrigins: List? = null, maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE, block: ServerSSESession.() -> Server, +) { + mcp( + enableDnsRebindingProtection = enableDnsRebindingProtection, + allowedHosts = allowedHosts, + allowedOrigins = allowedOrigins, + maxRequestBodySize = maxRequestBodySize, + transportDecorator = { it }, + block = block, + ) +} + +/** + * Registers MCP over SSE, decorating each transport before it is connected to its [Server]. + * + * The returned [Transport] must delegate lifecycle methods and callback registration to the supplied + * [SseServerTransport] so that Ktor request handling and the MCP session remain connected. + * + * @param enableDnsRebindingProtection whether to install [DnsRebindingProtection] on this route. + * @param allowedHosts hostnames allowed in the `Host` header. + * @param allowedOrigins origins allowed in the `Origin` header, compared by hostname only. + * @param maxRequestBodySize maximum allowed size, in bytes, of an incoming POST body. + * @param transportDecorator transforms each SDK-created transport before [Server.createSession] receives it. + * @param block factory block that creates the [Server] for the connection. + * @throws IllegalStateException if the [SSE] plugin is not installed. + */ +@KtorDsl +@Suppress("LongParameterList") +public fun Route.mcp( + enableDnsRebindingProtection: Boolean = true, + allowedHosts: List? = null, + allowedOrigins: List? = null, + maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE, + transportDecorator: (SseServerTransport) -> Transport, + block: ServerSSESession.() -> Server, ) { try { plugin(SSE) @@ -113,7 +192,7 @@ public fun Route.mcp( val transportManager = TransportManager() sse { - mcpSseEndpoint("", transportManager, maxRequestBodySize, block) + mcpSseEndpoint("", transportManager, maxRequestBodySize, transportDecorator, block) } post { @@ -149,12 +228,52 @@ public fun Application.mcp( allowedOrigins: List? = null, maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE, block: ServerSSESession.() -> Server, +) { + mcp( + enableDnsRebindingProtection = enableDnsRebindingProtection, + allowedHosts = allowedHosts, + allowedOrigins = allowedOrigins, + maxRequestBodySize = maxRequestBodySize, + transportDecorator = { it }, + block = block, + ) +} + +/** + * Configures MCP over SSE, decorating each transport before it is connected to its [Server]. + * + * The returned [Transport] must delegate lifecycle methods and callback registration to the supplied + * [SseServerTransport] so that Ktor request handling and the MCP session remain connected. + * + * @param enableDnsRebindingProtection whether to install [DnsRebindingProtection] on this route. + * @param allowedHosts hostnames allowed in the `Host` header. + * @param allowedOrigins origins allowed in the `Origin` header, compared by hostname only. + * @param maxRequestBodySize maximum allowed size, in bytes, of an incoming POST body. + * @param transportDecorator transforms each SDK-created transport before [Server.createSession] receives it. + * @param block factory block that creates the [Server] for the connection. + */ +@KtorDsl +@Suppress("LongParameterList") +public fun Application.mcp( + enableDnsRebindingProtection: Boolean = true, + allowedHosts: List? = null, + allowedOrigins: List? = null, + maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE, + transportDecorator: (SseServerTransport) -> Transport, + block: ServerSSESession.() -> Server, ) { installMcpContentNegotiation() install(SSE) routing { - mcp(enableDnsRebindingProtection, allowedHosts, allowedOrigins, maxRequestBodySize, block) + mcp( + enableDnsRebindingProtection = enableDnsRebindingProtection, + allowedHosts = allowedHosts, + allowedOrigins = allowedOrigins, + maxRequestBodySize = maxRequestBodySize, + transportDecorator = transportDecorator, + block = block, + ) } } @@ -166,6 +285,7 @@ private fun Application.mcpStreamableHttp( allowedOrigins: List?, configuration: StreamableHttpServerTransport.Configuration, sseHeartbeatConfig: (Heartbeat.() -> Unit)?, + transportDecorator: (StreamableHttpServerTransport) -> Transport, block: RoutingContext.() -> Server, ) { installMcpContentNegotiation() @@ -197,6 +317,7 @@ private fun Application.mcpStreamableHttp( val transport = streamableTransport( transportManager = transportManager, configuration = configuration, + transportDecorator = transportDecorator, block = block, ) ?: return@post @@ -245,6 +366,45 @@ public fun Application.mcpStreamableHttp( eventStore: EventStore? = null, sseHeartbeatConfig: (Heartbeat.() -> Unit)? = null, block: RoutingContext.() -> Server, +) { + mcpStreamableHttp( + path = path, + enableDnsRebindingProtection = enableDnsRebindingProtection, + allowedHosts = allowedHosts, + allowedOrigins = allowedOrigins, + eventStore = eventStore, + sseHeartbeatConfig = sseHeartbeatConfig, + transportDecorator = { it }, + block = block, + ) +} + +/** + * Configures MCP over Streamable HTTP, decorating each transport before it is connected to its [Server]. + * + * The returned [Transport] must delegate lifecycle methods and callback registration to the supplied + * [StreamableHttpServerTransport] so that Ktor request handling and the MCP session remain connected. + * + * @param path the base path for the MCP Streamable HTTP endpoint. + * @param enableDnsRebindingProtection whether DNS rebinding protection is enabled. + * @param allowedHosts hostnames allowed in the `Host` header. + * @param allowedOrigins origins allowed in the `Origin` header, compared by hostname only. + * @param eventStore optional storage for resumable event streams. + * @param sseHeartbeatConfig optional heartbeat configuration for SSE connections. + * @param transportDecorator transforms each SDK-created transport before [Server.createSession] receives it. + * @param block factory block that creates the [Server] for the connection. + */ +@KtorDsl +@Suppress("LongParameterList") +public fun Application.mcpStreamableHttp( + path: String = "/mcp", + enableDnsRebindingProtection: Boolean = true, + allowedHosts: List? = null, + allowedOrigins: List? = null, + eventStore: EventStore? = null, + sseHeartbeatConfig: (Heartbeat.() -> Unit)? = null, + transportDecorator: (StreamableHttpServerTransport) -> Transport, + block: RoutingContext.() -> Server, ) { mcpStreamableHttp( path = path, @@ -256,6 +416,7 @@ public fun Application.mcpStreamableHttp( enableJsonResponse = true, ), sseHeartbeatConfig = sseHeartbeatConfig, + transportDecorator = transportDecorator, block = block, ) } @@ -366,6 +527,7 @@ private suspend fun ServerSSESession.mcpSseEndpoint( postEndpoint: String, transportManager: TransportManager, maxRequestBodySize: Long, + transportDecorator: (SseServerTransport) -> Transport, block: ServerSSESession.() -> Server, ) { val transport = mcpSseTransport(postEndpoint, transportManager, maxRequestBodySize) @@ -377,7 +539,7 @@ private suspend fun ServerSSESession.mcpSseEndpoint( transportManager.removeTransport(transport.sessionId) } - server.createSession(transport) + server.createSession(transportDecorator(transport)) logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" } @@ -481,6 +643,7 @@ private suspend fun existingStreamableTransport( private suspend fun RoutingContext.streamableTransport( transportManager: TransportManager, configuration: StreamableHttpServerTransport.Configuration, + transportDecorator: (StreamableHttpServerTransport) -> Transport, block: RoutingContext.() -> Server, ): StreamableHttpServerTransport? { val sessionId = call.request.sessionId() @@ -506,7 +669,7 @@ private suspend fun RoutingContext.streamableTransport( transport.sessionId?.let { transportManager.removeTransport(it) } logger.info { "Server connection closed for sessionId: ${transport.sessionId}" } } - server.createSession(transport) + server.createSession(transportDecorator(transport)) return transport } diff --git a/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorTransportDecoratorTest.kt b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorTransportDecoratorTest.kt new file mode 100644 index 000000000..2587b8c11 --- /dev/null +++ b/kotlin-sdk-server/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorTransportDecoratorTest.kt @@ -0,0 +1,108 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.kotest.assertions.ktor.client.shouldHaveStatus +import io.kotest.matchers.shouldBe +import io.ktor.client.request.header +import io.ktor.client.request.post +import io.ktor.client.request.setBody +import io.ktor.http.ContentType +import io.ktor.http.HttpHeaders +import io.ktor.http.HttpStatusCode +import io.ktor.http.contentType +import io.ktor.server.application.install +import io.ktor.server.routing.routing +import io.ktor.server.sse.SSE +import io.ktor.server.testing.testApplication +import io.modelcontextprotocol.kotlin.sdk.shared.Transport +import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION +import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.Test + +class KtorTransportDecoratorTest : AbstractKtorExtensionsTest() { + + @Test + fun `Application mcp connects the decorated SSE transport`() = testApplication { + lateinit var decoratedTransport: StartTrackingTransport + + application { + mcp( + enableDnsRebindingProtection = false, + transportDecorator = { transport -> + StartTrackingTransport(transport).also { decoratedTransport = it } + }, + ) { testServer() } + } + + client.assertMcpEndpointsAt("/") + decoratedTransport.startCalls.get() shouldBe 1 + } + + @Test + fun `Route mcp with path connects the decorated SSE transport`() = testApplication { + lateinit var decoratedTransport: StartTrackingTransport + + application { + install(SSE) + routing { + mcp( + path = "/mcp", + enableDnsRebindingProtection = false, + transportDecorator = { transport -> + StartTrackingTransport(transport).also { decoratedTransport = it } + }, + ) { testServer() } + } + } + + client.assertMcpEndpointsAt("/mcp") + decoratedTransport.startCalls.get() shouldBe 1 + } + + @Test + fun `Streamable HTTP connects the decorated transport`() = testApplication { + lateinit var decoratedTransport: StartTrackingTransport + + application { + mcpStreamableHttp( + enableDnsRebindingProtection = false, + transportDecorator = { transport -> + StartTrackingTransport(transport).also { decoratedTransport = it } + }, + ) { testServer() } + } + + val response = client.post("/mcp") { + header( + HttpHeaders.Accept, + "${ContentType.Application.Json}, ${ContentType.Text.EventStream}", + ) + contentType(ContentType.Application.Json) + setBody( + """ + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "$LATEST_PROTOCOL_VERSION", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0.0"} + } + } + """.trimIndent(), + ) + } + + response.shouldHaveStatus(HttpStatusCode.OK) + decoratedTransport.startCalls.get() shouldBe 1 + } +} + +private class StartTrackingTransport(private val delegate: Transport) : Transport by delegate { + val startCalls: AtomicInteger = AtomicInteger() + + override suspend fun start() { + startCalls.incrementAndGet() + delegate.start() + } +}