Ktor extensions for streamable http (#504)
## Motivation and Context Adds Ktor extensions for streaming data over the server-side HTTP transport ## How Has This Been Tested? Integration test ## Breaking Changes - ## 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
This commit is contained in:
Pavel Gorgulov
GitHub
parent
c40fa16563
commit
da971aaee5
2 changed files with 356 additions and 16 deletions
|
|
@ -3,10 +3,15 @@ package io.modelcontextprotocol.kotlin.sdk.server
|
|||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.ktor.http.HttpStatusCode
|
||||
import io.ktor.server.application.Application
|
||||
import io.ktor.server.application.ApplicationCall
|
||||
import io.ktor.server.application.install
|
||||
import io.ktor.server.request.ApplicationRequest
|
||||
import io.ktor.server.request.header
|
||||
import io.ktor.server.response.respond
|
||||
import io.ktor.server.routing.Routing
|
||||
import io.ktor.server.routing.RoutingContext
|
||||
import io.ktor.server.routing.delete
|
||||
import io.ktor.server.routing.get
|
||||
import io.ktor.server.routing.post
|
||||
import io.ktor.server.routing.route
|
||||
import io.ktor.server.routing.routing
|
||||
|
|
@ -14,6 +19,8 @@ import io.ktor.server.sse.SSE
|
|||
import io.ktor.server.sse.ServerSSESession
|
||||
import io.ktor.server.sse.sse
|
||||
import io.ktor.utils.io.KtorDsl
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError
|
||||
import kotlinx.atomicfu.AtomicRef
|
||||
import kotlinx.atomicfu.atomic
|
||||
import kotlinx.atomicfu.update
|
||||
|
|
@ -23,13 +30,15 @@ import kotlinx.coroutines.awaitCancellation
|
|||
|
||||
private val logger = KotlinLogging.logger {}
|
||||
|
||||
internal class SseTransportManager(transports: Map<String, SseServerTransport> = emptyMap()) {
|
||||
private val transports: AtomicRef<PersistentMap<String, SseServerTransport>> = atomic(transports.toPersistentMap())
|
||||
internal class TransportManager(transports: Map<String, AbstractTransport> = emptyMap()) {
|
||||
private val transports: AtomicRef<PersistentMap<String, AbstractTransport>> = atomic(transports.toPersistentMap())
|
||||
|
||||
fun getTransport(sessionId: String): SseServerTransport? = transports.value[sessionId]
|
||||
fun hasTransport(sessionId: String): Boolean = transports.value.containsKey(sessionId)
|
||||
|
||||
fun addTransport(transport: SseServerTransport) {
|
||||
transports.update { it.put(transport.sessionId, transport) }
|
||||
fun getTransport(sessionId: String): AbstractTransport? = transports.value[sessionId]
|
||||
|
||||
fun addTransport(sessionId: String, transport: AbstractTransport) {
|
||||
transports.update { it.put(sessionId, transport) }
|
||||
}
|
||||
|
||||
fun removeTransport(sessionId: String) {
|
||||
|
|
@ -49,14 +58,14 @@ public fun Routing.mcp(path: String, block: ServerSSESession.() -> Server) {
|
|||
*/
|
||||
@KtorDsl
|
||||
public fun Routing.mcp(block: ServerSSESession.() -> Server) {
|
||||
val sseTransportManager = SseTransportManager()
|
||||
val transportManager = TransportManager()
|
||||
|
||||
sse {
|
||||
mcpSseEndpoint("", sseTransportManager, block)
|
||||
mcpSseEndpoint("", transportManager, block)
|
||||
}
|
||||
|
||||
post {
|
||||
mcpPostEndpoint(sseTransportManager)
|
||||
mcpPostEndpoint(transportManager)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -69,18 +78,98 @@ public fun Application.mcp(block: ServerSSESession.() -> Server) {
|
|||
}
|
||||
}
|
||||
|
||||
@KtorDsl
|
||||
public fun Application.mcpStreamableHttp(
|
||||
enableDnsRebindingProtection: Boolean = false,
|
||||
allowedHosts: List<String>? = null,
|
||||
allowedOrigins: List<String>? = null,
|
||||
eventStore: EventStore? = null,
|
||||
block: RoutingContext.() -> Server,
|
||||
) {
|
||||
install(SSE)
|
||||
|
||||
val transportManager = TransportManager()
|
||||
|
||||
routing {
|
||||
route("/mcp") {
|
||||
sse {
|
||||
val transport = existingStreamableTransport(call, transportManager) ?: return@sse
|
||||
transport.handleRequest(this, call)
|
||||
}
|
||||
|
||||
post {
|
||||
val transport = streamableTransport(
|
||||
transportManager = transportManager,
|
||||
enableDnsRebindingProtection = enableDnsRebindingProtection,
|
||||
allowedHosts = allowedHosts,
|
||||
allowedOrigins = allowedOrigins,
|
||||
eventStore = eventStore,
|
||||
block = block,
|
||||
)
|
||||
?: return@post
|
||||
|
||||
transport.handleRequest(null, call)
|
||||
}
|
||||
|
||||
delete {
|
||||
val transport = existingStreamableTransport(call, transportManager) ?: return@delete
|
||||
transport.handleRequest(null, call)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@KtorDsl
|
||||
public fun Application.mcpStatelessStreamableHttp(
|
||||
enableDnsRebindingProtection: Boolean = false,
|
||||
allowedHosts: List<String>? = null,
|
||||
allowedOrigins: List<String>? = null,
|
||||
eventStore: EventStore? = null,
|
||||
block: RoutingContext.() -> Server,
|
||||
) {
|
||||
install(SSE)
|
||||
|
||||
routing {
|
||||
route("/mcp") {
|
||||
post {
|
||||
mcpStatelessStreamableHttpEndpoint(
|
||||
enableDnsRebindingProtection = enableDnsRebindingProtection,
|
||||
allowedHosts = allowedHosts,
|
||||
allowedOrigins = allowedOrigins,
|
||||
eventStore = eventStore,
|
||||
block = block,
|
||||
)
|
||||
}
|
||||
get {
|
||||
call.reject(
|
||||
HttpStatusCode.MethodNotAllowed,
|
||||
RPCError.ErrorCode.CONNECTION_CLOSED,
|
||||
"Method not allowed.",
|
||||
)
|
||||
}
|
||||
delete {
|
||||
call.reject(
|
||||
HttpStatusCode.MethodNotAllowed,
|
||||
RPCError.ErrorCode.CONNECTION_CLOSED,
|
||||
"Method not allowed.",
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
internal suspend fun ServerSSESession.mcpSseEndpoint(
|
||||
postEndpoint: String,
|
||||
sseTransportManager: SseTransportManager,
|
||||
transportManager: TransportManager,
|
||||
block: ServerSSESession.() -> Server,
|
||||
) {
|
||||
val transport = mcpSseTransport(postEndpoint, sseTransportManager)
|
||||
val transport = mcpSseTransport(postEndpoint, transportManager)
|
||||
|
||||
val server = block()
|
||||
|
||||
server.onClose {
|
||||
logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
|
||||
sseTransportManager.removeTransport(transport.sessionId)
|
||||
transportManager.removeTransport(transport.sessionId)
|
||||
}
|
||||
|
||||
server.createSession(transport)
|
||||
|
|
@ -92,16 +181,41 @@ internal suspend fun ServerSSESession.mcpSseEndpoint(
|
|||
|
||||
internal fun ServerSSESession.mcpSseTransport(
|
||||
postEndpoint: String,
|
||||
sseTransportManager: SseTransportManager,
|
||||
transportManager: TransportManager,
|
||||
): SseServerTransport {
|
||||
val transport = SseServerTransport(postEndpoint, this)
|
||||
sseTransportManager.addTransport(transport)
|
||||
transportManager.addTransport(transport.sessionId, transport)
|
||||
logger.info { "New SSE connection established and stored with sessionId: ${transport.sessionId}" }
|
||||
|
||||
return transport
|
||||
}
|
||||
|
||||
internal suspend fun RoutingContext.mcpPostEndpoint(sseTransportManager: SseTransportManager) {
|
||||
internal suspend fun RoutingContext.mcpStatelessStreamableHttpEndpoint(
|
||||
enableDnsRebindingProtection: Boolean = false,
|
||||
allowedHosts: List<String>? = null,
|
||||
allowedOrigins: List<String>? = null,
|
||||
eventStore: EventStore? = null,
|
||||
block: RoutingContext.() -> Server,
|
||||
) {
|
||||
val transport = StreamableHttpServerTransport(
|
||||
enableDnsRebindingProtection = enableDnsRebindingProtection,
|
||||
allowedHosts = allowedHosts,
|
||||
allowedOrigins = allowedOrigins,
|
||||
eventStore = eventStore,
|
||||
enableJsonResponse = true,
|
||||
).also { it.setSessionIdGenerator(null) }
|
||||
|
||||
logger.info { "New stateless StreamableHttp connection established without sessionId" }
|
||||
|
||||
val server = block()
|
||||
server.onClose { logger.info { "Server connection closed without sessionId" } }
|
||||
server.createSession(transport)
|
||||
|
||||
transport.handleRequest(null, this.call)
|
||||
logger.debug { "Server connected to transport without sessionId" }
|
||||
}
|
||||
|
||||
internal suspend fun RoutingContext.mcpPostEndpoint(transportManager: TransportManager) {
|
||||
val sessionId: String = call.request.queryParameters["sessionId"] ?: run {
|
||||
call.respond(HttpStatusCode.BadRequest, "sessionId query parameter is not provided")
|
||||
return
|
||||
|
|
@ -109,7 +223,7 @@ internal suspend fun RoutingContext.mcpPostEndpoint(sseTransportManager: SseTran
|
|||
|
||||
logger.debug { "Received message for sessionId: $sessionId" }
|
||||
|
||||
val transport = sseTransportManager.getTransport(sessionId)
|
||||
val transport = transportManager.getTransport(sessionId) as SseServerTransport?
|
||||
if (transport == null) {
|
||||
logger.warn { "Session not found for sessionId: $sessionId" }
|
||||
call.respond(HttpStatusCode.NotFound, "Session not found")
|
||||
|
|
@ -119,3 +233,74 @@ internal suspend fun RoutingContext.mcpPostEndpoint(sseTransportManager: SseTran
|
|||
transport.handlePostMessage(call)
|
||||
logger.trace { "Message handled for sessionId: $sessionId" }
|
||||
}
|
||||
|
||||
private fun ApplicationRequest.sessionId(): String? = header(MCP_SESSION_ID_HEADER)
|
||||
|
||||
private suspend fun existingStreamableTransport(
|
||||
call: ApplicationCall,
|
||||
transportManager: TransportManager,
|
||||
): StreamableHttpServerTransport? {
|
||||
val sessionId = call.request.sessionId()
|
||||
if (sessionId.isNullOrEmpty()) {
|
||||
call.reject(
|
||||
HttpStatusCode.BadRequest,
|
||||
RPCError.ErrorCode.CONNECTION_CLOSED,
|
||||
"Bad Request: No valid session ID provided",
|
||||
)
|
||||
return null
|
||||
}
|
||||
|
||||
val transport = transportManager.getTransport(sessionId) as? StreamableHttpServerTransport
|
||||
if (transport == null) {
|
||||
call.reject(
|
||||
HttpStatusCode.NotFound,
|
||||
RPCError.ErrorCode.CONNECTION_CLOSED,
|
||||
"Session not found",
|
||||
)
|
||||
return null
|
||||
}
|
||||
|
||||
return transport
|
||||
}
|
||||
|
||||
private suspend fun RoutingContext.streamableTransport(
|
||||
transportManager: TransportManager,
|
||||
enableDnsRebindingProtection: Boolean,
|
||||
allowedHosts: List<String>?,
|
||||
allowedOrigins: List<String>?,
|
||||
eventStore: EventStore?,
|
||||
block: RoutingContext.() -> Server,
|
||||
): StreamableHttpServerTransport? {
|
||||
val sessionId = call.request.sessionId()
|
||||
if (sessionId != null) {
|
||||
val transport = transportManager.getTransport(sessionId) as? StreamableHttpServerTransport
|
||||
return transport ?: existingStreamableTransport(call, transportManager)
|
||||
}
|
||||
|
||||
val transport = StreamableHttpServerTransport(
|
||||
enableDnsRebindingProtection = enableDnsRebindingProtection,
|
||||
allowedHosts = allowedHosts,
|
||||
allowedOrigins = allowedOrigins,
|
||||
eventStore = eventStore,
|
||||
enableJsonResponse = true,
|
||||
)
|
||||
|
||||
transport.setOnSessionInitialized { initializedSessionId ->
|
||||
transportManager.addTransport(initializedSessionId, transport)
|
||||
logger.info { "New StreamableHttp connection established and stored with sessionId: $initializedSessionId" }
|
||||
}
|
||||
|
||||
transport.setOnSessionClosed { closedSession ->
|
||||
transportManager.removeTransport(closedSession)
|
||||
logger.info { "Closed StreamableHttp connection and removed sessionId: $closedSession" }
|
||||
}
|
||||
|
||||
val server = block()
|
||||
server.onClose {
|
||||
transport.sessionId?.let { transportManager.removeTransport(it) }
|
||||
logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
|
||||
}
|
||||
server.createSession(transport)
|
||||
|
||||
return transport
|
||||
}
|
||||
|
|
|
|||
|
|
@ -254,7 +254,9 @@ public class StreamableHttpServerTransport(
|
|||
*/
|
||||
public suspend fun handlePostRequest(session: ServerSSESession?, call: ApplicationCall) {
|
||||
try {
|
||||
if (!enableJsonResponse && session == null) error("Server session can't be null with json response")
|
||||
if (!enableJsonResponse && session == null) {
|
||||
error("Server session can't be null for SSE responses")
|
||||
}
|
||||
|
||||
val acceptHeader = call.request.header(HttpHeaders.Accept)
|
||||
val isAcceptEventStream = acceptHeader.accepts(ContentType.Text.EventStream)
|
||||
|
|
|
|||
Loading…
Reference in a new issue