Skip to content

feat(server): allow decorating Ktor transports #953

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ECD5A wants to merge 1 commit into
base: main
Choose a base branch
from
+284 −5
Open

feat(server): allow decorating Ktor transports #953

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions kotlin-sdk-server/api/kotlin-sdk-server.api
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
173 changes: 168 additions & 5 deletions ...-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/KtorServer.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -62,9 +63,53 @@ public fun Route.mcp(
allowedOrigins: List<String>? = 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<String>? = null,
allowedOrigins: List<String>? = 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,
)
}
}

Expand Down Expand Up @@ -96,6 +141,40 @@ public fun Route.mcp(
allowedOrigins: List<String>? = 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<String>? = null,
allowedOrigins: List<String>? = null,
maxRequestBodySize: Long = DEFAULT_MAX_REQUEST_BODY_SIZE,
transportDecorator: (SseServerTransport) -> Transport,
block: ServerSSESession.() -> Server,
) {
try {
plugin(SSE)
Expand All @@ -113,7 +192,7 @@ public fun Route.mcp(
val transportManager = TransportManager<SseServerTransport>()

sse {
mcpSseEndpoint("", transportManager, maxRequestBodySize, block)
mcpSseEndpoint("", transportManager, maxRequestBodySize, transportDecorator, block)
}

post {
Expand Down Expand Up @@ -149,12 +228,52 @@ public fun Application.mcp(
allowedOrigins: List<String>? = 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<String>? = null,
allowedOrigins: List<String>? = 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,
)
}
}

Expand All @@ -166,6 +285,7 @@ private fun Application.mcpStreamableHttp(
allowedOrigins: List<String>?,
configuration: StreamableHttpServerTransport.Configuration,
sseHeartbeatConfig: (Heartbeat.() -> Unit)?,
transportDecorator: (StreamableHttpServerTransport) -> Transport,
block: RoutingContext.() -> Server,
) {
installMcpContentNegotiation()
Expand Down Expand Up @@ -197,6 +317,7 @@ private fun Application.mcpStreamableHttp(
val transport = streamableTransport(
transportManager = transportManager,
configuration = configuration,
transportDecorator = transportDecorator,
block = block,
) ?: return@post

Expand Down Expand Up @@ -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<String>? = null,
allowedOrigins: List<String>? = null,
eventStore: EventStore? = null,
sseHeartbeatConfig: (Heartbeat.() -> Unit)? = null,
transportDecorator: (StreamableHttpServerTransport) -> Transport,
block: RoutingContext.() -> Server,
) {
mcpStreamableHttp(
path = path,
Expand All @@ -256,6 +416,7 @@ public fun Application.mcpStreamableHttp(
enableJsonResponse = true,
),
sseHeartbeatConfig = sseHeartbeatConfig,
transportDecorator = transportDecorator,
block = block,
)
}
Expand Down Expand Up @@ -366,6 +527,7 @@ private suspend fun ServerSSESession.mcpSseEndpoint(
postEndpoint: String,
transportManager: TransportManager<SseServerTransport>,
maxRequestBodySize: Long,
transportDecorator: (SseServerTransport) -> Transport,
block: ServerSSESession.() -> Server,
) {
val transport = mcpSseTransport(postEndpoint, transportManager, maxRequestBodySize)
Expand All @@ -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}" }

Expand Down Expand Up @@ -481,6 +643,7 @@ private suspend fun existingStreamableTransport(
private suspend fun RoutingContext.streamableTransport(
transportManager: TransportManager<StreamableHttpServerTransport>,
configuration: StreamableHttpServerTransport.Configuration,
transportDecorator: (StreamableHttpServerTransport) -> Transport,
block: RoutingContext.() -> Server,
): StreamableHttpServerTransport? {
val sessionId = call.request.sessionId()
Expand All @@ -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
}
Expand Down
Loading