diff --git a/gradle.properties b/gradle.properties index c4128f5..9abdd0f 100644 --- a/gradle.properties +++ b/gradle.properties @@ -14,4 +14,4 @@ kotlin.mpp.enableCInteropCommonization=true kotlin.native.ignoreDisabledTargets=true group=io.modelcontextprotocol -version=0.8.3 +version=0.8.4 diff --git a/integration-test/build.gradle.kts b/integration-test/build.gradle.kts index d5b8ef8..35ca4af 100644 --- a/integration-test/build.gradle.kts +++ b/integration-test/build.gradle.kts @@ -19,6 +19,8 @@ kotlin { implementation(libs.kotlinx.coroutines.test) implementation(libs.ktor.client.cio) implementation(libs.ktor.server.cio) + implementation(libs.ktor.server.content.negotiation) + implementation(libs.ktor.serialization) implementation(libs.ktor.server.websockets) implementation(libs.ktor.server.test.host) } 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 new file mode 100644 index 0000000..0f6646e --- /dev/null +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/AbstractStreamableHttpIntegrationTest.kt @@ -0,0 +1,82 @@ +package io.modelcontextprotocol.kotlin.sdk.integration.streamablehttp + +import io.ktor.serialization.kotlinx.json.json +import io.ktor.server.application.install +import io.ktor.server.cio.CIOApplicationEngine +import io.ktor.server.engine.EmbeddedServer +import io.ktor.server.engine.embeddedServer +import io.ktor.server.plugins.contentnegotiation.ContentNegotiation +import io.modelcontextprotocol.kotlin.sdk.server.Server +import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions +import io.modelcontextprotocol.kotlin.sdk.server.mcpStreamableHttp +import io.modelcontextprotocol.kotlin.sdk.types.GetPromptResult +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.McpJson +import io.modelcontextprotocol.kotlin.sdk.types.PromptArgument +import io.modelcontextprotocol.kotlin.sdk.types.PromptMessage +import io.modelcontextprotocol.kotlin.sdk.types.Role +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import io.ktor.server.cio.CIO as ServerCIO + +open class AbstractStreamableHttpIntegrationTest { + + suspend fun initTestServer(name: String? = null): StreamableHttpTestServer { + val mcpServer = Server( + Implementation(name = name ?: DEFAULT_SERVER_NAME, version = VERSION), + ServerOptions( + capabilities = ServerCapabilities(prompts = ServerCapabilities.Prompts(listChanged = true)), + ), + ) { + addPrompt( + name = "prompt", + description = "Prompt description", + arguments = listOf( + PromptArgument( + name = "client", + description = "Client name who requested a prompt", + required = true, + ), + ), + ) { request -> + GetPromptResult( + messages = listOf( + PromptMessage( + role = Role.User, + content = TextContent("Prompt for client ${request.params.arguments?.get("client")}"), + ), + ), + description = "Prompt for ${request.params.name}", + ) + } + } + + val ktorServer = embeddedServer( + ServerCIO, + host = URL, + port = PORT, + ) { + install(ContentNegotiation) { + json(McpJson) + } + mcpStreamableHttp { mcpServer } + } + + return StreamableHttpTestServer( + mcpServer = mcpServer, + ktorServer = ktorServer.startSuspend(wait = false), + ) + } + + data class StreamableHttpTestServer( + val mcpServer: Server, + val ktorServer: EmbeddedServer, + ) + + companion object { + internal const val DEFAULT_SERVER_NAME = "streamable-http-test-server" + internal const val VERSION = "1.0.0" + internal const val URL = "127.0.0.1" + internal const val PORT = 0 + } +} diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpIntegrationTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpIntegrationTest.kt new file mode 100644 index 0000000..b64df4a --- /dev/null +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/streamablehttp/StreamableHttpIntegrationTest.kt @@ -0,0 +1,65 @@ +package io.modelcontextprotocol.kotlin.sdk.integration.streamablehttp + +import io.ktor.client.HttpClient +import io.ktor.client.plugins.sse.SSE +import io.modelcontextprotocol.kotlin.sdk.client.Client +import io.modelcontextprotocol.kotlin.sdk.client.StreamableHttpClientTransport +import io.modelcontextprotocol.kotlin.sdk.types.GetPromptRequest +import io.modelcontextprotocol.kotlin.sdk.types.GetPromptRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.Implementation +import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import io.modelcontextprotocol.kotlin.test.utils.actualPort +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertTrue +import io.ktor.client.engine.cio.CIO as ClientCIO + +class StreamableHttpIntegrationTest : AbstractStreamableHttpIntegrationTest() { + + @Test + fun `client should receive responses via POST json and server notifications via GET sse`() = + runBlocking(Dispatchers.IO) { + var server: StreamableHttpTestServer? = null + var client: Client? = null + var httpClient: HttpClient? = null + + try { + server = initTestServer() + val port = server.ktorServer.actualPort() + + httpClient = HttpClient(ClientCIO) { install(SSE) } + val transport = StreamableHttpClientTransport( + client = httpClient, + url = "http://$URL:$port/mcp", + ).apply { + protocolVersion = LATEST_PROTOCOL_VERSION + } + + client = Client(Implementation("streamable-http-client", VERSION)) + client.connect(transport) + + val prompt = getPrompt(client, "Client A") + assertTrue { "Client A" in prompt } + } finally { + client?.close() + httpClient?.close() + server?.ktorServer?.stopSuspend(1000, 2000) + } + } + + private suspend fun getPrompt(client: Client, clientName: String): String { + val response = client.getPrompt( + GetPromptRequest( + GetPromptRequestParams( + name = "prompt", + arguments = mapOf("client" to clientName), + ), + ), + ) + + return (response.messages.first().content as? TextContent)?.text + ?: error("Failed to receive prompt for Client $clientName") + } +} diff --git a/kotlin-sdk-core/api/kotlin-sdk-core.api b/kotlin-sdk-core/api/kotlin-sdk-core.api index 67acbdd..86d0bbb 100644 --- a/kotlin-sdk-core/api/kotlin-sdk-core.api +++ b/kotlin-sdk-core/api/kotlin-sdk-core.api @@ -4492,13 +4492,15 @@ public final class io/modelcontextprotocol/kotlin/sdk/types/ToolListChangedNotif public final class io/modelcontextprotocol/kotlin/sdk/types/ToolSchema { public static final field Companion Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema$Companion; public fun ()V - public fun (Lkotlinx/serialization/json/JsonObject;Ljava/util/List;)V - public synthetic fun (Lkotlinx/serialization/json/JsonObject;Ljava/util/List;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lkotlinx/serialization/json/JsonObject;Ljava/util/List;Lkotlinx/serialization/json/JsonObject;)V + public synthetic fun (Lkotlinx/serialization/json/JsonObject;Ljava/util/List;Lkotlinx/serialization/json/JsonObject;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Lkotlinx/serialization/json/JsonObject; public final fun component2 ()Ljava/util/List; - public final fun copy (Lkotlinx/serialization/json/JsonObject;Ljava/util/List;)Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema; - public static synthetic fun copy$default (Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema;Lkotlinx/serialization/json/JsonObject;Ljava/util/List;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema; + public final fun component3 ()Lkotlinx/serialization/json/JsonObject; + public final fun copy (Lkotlinx/serialization/json/JsonObject;Ljava/util/List;Lkotlinx/serialization/json/JsonObject;)Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema; + public static synthetic fun copy$default (Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema;Lkotlinx/serialization/json/JsonObject;Ljava/util/List;Lkotlinx/serialization/json/JsonObject;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/types/ToolSchema; public fun equals (Ljava/lang/Object;)Z + public final fun getDefs ()Lkotlinx/serialization/json/JsonObject; public final fun getProperties ()Lkotlinx/serialization/json/JsonObject; public final fun getRequired ()Ljava/util/List; public final fun getType ()Ljava/lang/String; diff --git a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/tools.kt b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/tools.kt index 6cba208..5d975c6 100644 --- a/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/tools.kt +++ b/kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/types/tools.kt @@ -78,9 +78,15 @@ public data class Tool( * @property type Always "object" for tool schemas. * @property properties Optional map of property names to their schema definitions. * @property required Optional list of property names that are required. + * @property defs Optional schema definitions available to references in [properties]. Serialized as `$defs`. */ @Serializable -public data class ToolSchema(val properties: JsonObject? = null, val required: List? = null) { +public data class ToolSchema( + val properties: JsonObject? = null, + val required: List? = null, + @SerialName("\$defs") + val defs: JsonObject? = null, +) { @OptIn(ExperimentalSerializationApi::class) @EncodeDefault val type: String = "object" diff --git a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/ToolsTest.kt b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/ToolsTest.kt index a662e44..447df4f 100644 --- a/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/ToolsTest.kt +++ b/kotlin-sdk-core/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/types/ToolsTest.kt @@ -152,6 +152,57 @@ class ToolsTest { assertEquals("language", tool.meta?.get("category")?.jsonPrimitive?.content) } + @Test + fun `should serialize ToolSchema with defs`() { + val tool = toolWithDefs() + + val json = McpJson.encodeToString(tool) + + json shouldEqualJson toolWithDefsJson() + } + + @Test + fun `should deserialize ToolSchema with defs`() { + val json = $$""" + { + "name": "create-page", + "inputSchema": { + "type": "object", + "$defs": { + "parentRequest": { + "type": "object", + "properties": { + "page_id": { + "type": "string" + } + } + } + }, + "properties": { + "parent": { + "$ref": "#/$defs/parentRequest" + } + }, + "required": ["parent"] + } + } + """.trimIndent() + + val tool = McpJson.decodeFromString(json) + + val schema = tool.inputSchema + val defs = schema.defs + assertNotNull(defs) + val parentRequest = defs["parentRequest"]?.jsonObject + assertNotNull(parentRequest) + assertEquals("object", parentRequest["type"]?.jsonPrimitive?.content) + assertEquals( + $$"#/$defs/parentRequest", + schema.properties?.get("parent")?.jsonObject?.get($$"$ref")?.jsonPrimitive?.content, + ) + assertEquals(listOf("parent"), schema.required) + } + @Test fun `should serialize CallToolRequest with arguments`() { val request = CallToolRequest( @@ -520,4 +571,65 @@ class ToolsTest { } } """.trimIndent() + + private fun toolWithDefs(): Tool = Tool( + name = "create-page", + inputSchema = toolSchemaWithDefs(), + ) + + private fun toolSchemaWithDefs(): ToolSchema = ToolSchema( + properties = buildJsonObject { + put( + "parent", + buildJsonObject { + put($$"$ref", $$"#/$defs/parentRequest") + }, + ) + }, + required = listOf("parent"), + defs = buildJsonObject { + put( + "parentRequest", + buildJsonObject { + put("type", "object") + put( + "properties", + buildJsonObject { + put( + "page_id", + buildJsonObject { + put("type", "string") + }, + ) + }, + ) + }, + ) + }, + ) + + private fun toolWithDefsJson(): String = $$""" + { + "name": "create-page", + "inputSchema": { + "type": "object", + "$defs": { + "parentRequest": { + "type": "object", + "properties": { + "page_id": { + "type": "string" + } + } + } + }, + "properties": { + "parent": { + "$ref": "#/$defs/parentRequest" + } + }, + "required": ["parent"] + } + } + """.trimIndent() } diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index 53c971b..f0eb5c1 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -8,6 +8,10 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/KtorServerKt { public static final fun mcp (Lio/ktor/server/application/Application;Lkotlin/jvm/functions/Function1;)V public static final fun mcp (Lio/ktor/server/routing/Routing;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public static final fun mcp (Lio/ktor/server/routing/Routing;Lkotlin/jvm/functions/Function1;)V + public static final fun mcpStatelessStreamableHttp (Lio/ktor/server/application/Application;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;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;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;ZLjava/util/List;Ljava/util/List;Lio/modelcontextprotocol/kotlin/sdk/server/EventStore;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 2a2f950..8bb09a8 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 @@ -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 = emptyMap()) { - private val transports: AtomicRef> = atomic(transports.toPersistentMap()) +internal class TransportManager(transports: Map = emptyMap()) { + private val transports: AtomicRef> = 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? = null, + allowedOrigins: List? = 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? = null, + allowedOrigins: List? = 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? = null, + allowedOrigins: List? = 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?, + allowedOrigins: List?, + 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 +} diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt index 56bcead..e4ccf71 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StreamableHttpServerTransport.kt @@ -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)