modularize sdk (#208)
* Modularize SDK: migrated client and server implementations to separate modules. Removed old configuration files and reorganized the source code structure * add support native targets and update js/wasm configurations * Refactor sample projects: add samples as composite builds, replace hardcoded dependencies with version catalog, remove obsolete gradle wrapper files * Update .gitignore: exclude SWE agent directories (.claude/ and .junie/) * Refine CodeQL workflow: update Kotlin build steps to target individual modules separately instead of building all at once
This commit is contained in:
Pavel Gorgulov
GitHub
parent
641df74f5d
commit
8698b96dd2
7 changed files with 979 additions and 2003 deletions
|
|
@ -0,0 +1,117 @@
|
|||
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.install
|
||||
import io.ktor.server.response.respond
|
||||
import io.ktor.server.routing.Routing
|
||||
import io.ktor.server.routing.RoutingContext
|
||||
import io.ktor.server.routing.post
|
||||
import io.ktor.server.routing.route
|
||||
import io.ktor.server.routing.routing
|
||||
import io.ktor.server.sse.SSE
|
||||
import io.ktor.server.sse.ServerSSESession
|
||||
import io.ktor.server.sse.sse
|
||||
import io.ktor.util.collections.ConcurrentMap
|
||||
import io.ktor.utils.io.KtorDsl
|
||||
|
||||
private val logger = KotlinLogging.logger {}
|
||||
|
||||
@KtorDsl
|
||||
public fun Routing.mcp(path: String, block: () -> Server) {
|
||||
route(path) {
|
||||
mcp(block)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the Ktor Application to handle Model Context Protocol (MCP) over Server-Sent Events (SSE).
|
||||
*/
|
||||
@KtorDsl
|
||||
public fun Routing.mcp(block: () -> Server) {
|
||||
val transports = ConcurrentMap<String, SseServerTransport>()
|
||||
|
||||
sse {
|
||||
mcpSseEndpoint("", transports, block)
|
||||
}
|
||||
|
||||
post {
|
||||
mcpPostEndpoint(transports)
|
||||
}
|
||||
}
|
||||
|
||||
@Suppress("FunctionName")
|
||||
@Deprecated("Use mcp() instead", ReplaceWith("mcp(block)"), DeprecationLevel.WARNING)
|
||||
public fun Application.MCP(block: () -> Server) {
|
||||
mcp(block)
|
||||
}
|
||||
|
||||
@KtorDsl
|
||||
public fun Application.mcp(block: () -> Server) {
|
||||
val transports = ConcurrentMap<String, SseServerTransport>()
|
||||
|
||||
install(SSE)
|
||||
|
||||
routing {
|
||||
sse("/sse") {
|
||||
mcpSseEndpoint("/message", transports, block)
|
||||
}
|
||||
|
||||
post("/message") {
|
||||
mcpPostEndpoint(transports)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun ServerSSESession.mcpSseEndpoint(
|
||||
postEndpoint: String,
|
||||
transports: ConcurrentMap<String, SseServerTransport>,
|
||||
block: () -> Server,
|
||||
) {
|
||||
val transport = mcpSseTransport(postEndpoint, transports)
|
||||
|
||||
val server = block()
|
||||
|
||||
server.onClose {
|
||||
logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
|
||||
transports.remove(transport.sessionId)
|
||||
}
|
||||
|
||||
server.connect(transport)
|
||||
logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" }
|
||||
}
|
||||
|
||||
internal fun ServerSSESession.mcpSseTransport(
|
||||
postEndpoint: String,
|
||||
transports: ConcurrentMap<String, SseServerTransport>,
|
||||
): SseServerTransport {
|
||||
val transport = SseServerTransport(postEndpoint, this)
|
||||
transports[transport.sessionId] = transport
|
||||
|
||||
logger.info { "New SSE connection established and stored with sessionId: ${transport.sessionId}" }
|
||||
|
||||
return transport
|
||||
}
|
||||
|
||||
internal suspend fun RoutingContext.mcpPostEndpoint(
|
||||
transports: ConcurrentMap<String, SseServerTransport>,
|
||||
) {
|
||||
val sessionId: String = call.request.queryParameters["sessionId"]
|
||||
?: run {
|
||||
call.respond(HttpStatusCode.BadRequest, "sessionId query parameter is not provided")
|
||||
return
|
||||
}
|
||||
|
||||
logger.debug { "Received message for sessionId: $sessionId" }
|
||||
|
||||
val transport = transports[sessionId]
|
||||
if (transport == null) {
|
||||
logger.warn { "Session not found for sessionId: $sessionId" }
|
||||
call.respond(HttpStatusCode.NotFound, "Session not found")
|
||||
return
|
||||
}
|
||||
|
||||
transport.handlePostMessage(call)
|
||||
logger.trace { "Message handled for sessionId: $sessionId" }
|
||||
}
|
||||
|
|
@ -0,0 +1,127 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.ktor.http.ContentType
|
||||
import io.ktor.http.HttpStatusCode
|
||||
import io.ktor.http.encodeURLPath
|
||||
import io.ktor.server.application.ApplicationCall
|
||||
import io.ktor.server.request.contentType
|
||||
import io.ktor.server.request.receiveText
|
||||
import io.ktor.server.response.respondText
|
||||
import io.ktor.server.sse.ServerSSESession
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.McpJson
|
||||
import kotlinx.coroutines.job
|
||||
import kotlin.concurrent.atomics.AtomicBoolean
|
||||
import kotlin.concurrent.atomics.ExperimentalAtomicApi
|
||||
import kotlin.uuid.ExperimentalUuidApi
|
||||
import kotlin.uuid.Uuid
|
||||
|
||||
internal const val SESSION_ID_PARAM = "sessionId"
|
||||
|
||||
@Deprecated("Use SseServerTransport instead", ReplaceWith("SseServerTransport"), DeprecationLevel.WARNING)
|
||||
public typealias SSEServerTransport = SseServerTransport
|
||||
|
||||
/**
|
||||
* Server transport for SSE: this will send messages over an SSE connection and receive messages from HTTP POST requests.
|
||||
*
|
||||
* Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL identified by `_endpoint`.
|
||||
*/
|
||||
@OptIn(ExperimentalAtomicApi::class)
|
||||
public class SseServerTransport(
|
||||
private val endpoint: String,
|
||||
private val session: ServerSSESession,
|
||||
) : AbstractTransport() {
|
||||
private val initialized: AtomicBoolean = AtomicBoolean(false)
|
||||
|
||||
@OptIn(ExperimentalUuidApi::class)
|
||||
public val sessionId: String = Uuid.random().toString()
|
||||
|
||||
/**
|
||||
* Handles the initial SSE connection request.
|
||||
*
|
||||
* This should be called when a GET request is made to establish the SSE stream.
|
||||
*/
|
||||
override suspend fun start() {
|
||||
if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
|
||||
error("SSEServerTransport already started! If using Server class, note that connect() calls start() automatically.")
|
||||
}
|
||||
|
||||
// Send the endpoint event
|
||||
session.send(
|
||||
event = "endpoint",
|
||||
data = "${endpoint.encodeURLPath()}?$SESSION_ID_PARAM=${sessionId}",
|
||||
)
|
||||
|
||||
try {
|
||||
session.coroutineContext.job.join()
|
||||
} finally {
|
||||
_onClose.invoke()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles incoming POST messages.
|
||||
*
|
||||
* This should be called when a POST request is made to send a message to the server.
|
||||
*/
|
||||
public suspend fun handlePostMessage(call: ApplicationCall) {
|
||||
if (!initialized.load()) {
|
||||
val message = "SSE connection not established"
|
||||
call.respondText(message, status = HttpStatusCode.InternalServerError)
|
||||
_onError.invoke(IllegalStateException(message))
|
||||
}
|
||||
|
||||
val body = try {
|
||||
val ct = call.request.contentType()
|
||||
if (ct != ContentType.Application.Json) {
|
||||
error("Unsupported content-type: $ct")
|
||||
}
|
||||
|
||||
call.receiveText()
|
||||
} catch (e: Exception) {
|
||||
call.respondText("Invalid message: ${e.message}", status = HttpStatusCode.BadRequest)
|
||||
_onError.invoke(e)
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
handleMessage(body)
|
||||
} catch (e: Exception) {
|
||||
call.respondText("Error handling message $body: ${e.message}", status = HttpStatusCode.BadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
call.respondText("Accepted", status = HttpStatusCode.Accepted)
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle a client message, regardless of how it arrived.
|
||||
* This can be used to inform the server of messages that arrive via a means different from HTTP POST.
|
||||
*/
|
||||
public suspend fun handleMessage(message: String) {
|
||||
try {
|
||||
val parsedMessage = McpJson.decodeFromString<JSONRPCMessage>(message)
|
||||
_onMessage.invoke(parsedMessage)
|
||||
} catch (e: Exception) {
|
||||
_onError.invoke(e)
|
||||
throw e
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
session.close()
|
||||
_onClose.invoke()
|
||||
}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage) {
|
||||
if (!initialized.load()) {
|
||||
error("Not connected")
|
||||
}
|
||||
|
||||
session.send(
|
||||
event = "message",
|
||||
data = McpJson.encodeToString(message),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,834 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.modelcontextprotocol.kotlin.sdk.CallToolRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.CallToolResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ClientCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateElicitationRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateElicitationRequest.RequestedSchema
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateElicitationResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateMessageRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.CreateMessageResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.EmptyJsonObject
|
||||
import io.modelcontextprotocol.kotlin.sdk.EmptyRequestResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializeRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializeResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.LATEST_PROTOCOL_VERSION
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListPromptsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListPromptsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListResourcesRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListResourcesResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListRootsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListRootsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListToolsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ListToolsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.LoggingMessageNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.PingRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.Prompt
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptArgument
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptListChangedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.ReadResourceRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.Resource
|
||||
import io.modelcontextprotocol.kotlin.sdk.ResourceListChangedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.ResourceUpdatedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.SUPPORTED_PROTOCOL_VERSIONS
|
||||
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.Tool
|
||||
import io.modelcontextprotocol.kotlin.sdk.ToolAnnotations
|
||||
import io.modelcontextprotocol.kotlin.sdk.ToolListChangedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.Protocol
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.ProtocolOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
|
||||
import kotlinx.atomicfu.atomic
|
||||
import kotlinx.atomicfu.getAndUpdate
|
||||
import kotlinx.atomicfu.update
|
||||
import kotlinx.collections.immutable.minus
|
||||
import kotlinx.collections.immutable.persistentMapOf
|
||||
import kotlinx.collections.immutable.toPersistentSet
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
|
||||
private val logger = KotlinLogging.logger {}
|
||||
|
||||
/**
|
||||
* Configuration options for the MCP server.
|
||||
*
|
||||
* @property capabilities The capabilities this server supports.
|
||||
* @property enforceStrictCapabilities Whether to strictly enforce capabilities when interacting with clients.
|
||||
*/
|
||||
public class ServerOptions(
|
||||
public val capabilities: ServerCapabilities,
|
||||
enforceStrictCapabilities: Boolean = true,
|
||||
) : ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities)
|
||||
|
||||
/**
|
||||
* An MCP server on top of a pluggable transport.
|
||||
*
|
||||
* This server automatically responds to the initialization flow as initiated by the client.
|
||||
* You can register tools, prompts, and resources using [addTool], [addPrompt], and [addResource].
|
||||
* The server will then automatically handle listing and retrieval requests from the client.
|
||||
*
|
||||
* @param serverInfo Information about this server implementation (name, version).
|
||||
* @param options Configuration options for the server.
|
||||
*/
|
||||
public open class Server(
|
||||
private val serverInfo: Implementation,
|
||||
options: ServerOptions,
|
||||
) : Protocol(options) {
|
||||
private var _onInitialized: (() -> Unit) = {}
|
||||
private var _onClose: () -> Unit = {}
|
||||
|
||||
/**
|
||||
* The client's reported capabilities after initialization.
|
||||
*/
|
||||
public var clientCapabilities: ClientCapabilities? = null
|
||||
private set
|
||||
|
||||
/**
|
||||
* The client's version information after initialization.
|
||||
*/
|
||||
public var clientVersion: Implementation? = null
|
||||
private set
|
||||
|
||||
private val capabilities: ServerCapabilities = options.capabilities
|
||||
|
||||
private val _tools = atomic(persistentMapOf<String, RegisteredTool>())
|
||||
private val _prompts = atomic(persistentMapOf<String, RegisteredPrompt>())
|
||||
private val _resources = atomic(persistentMapOf<String, RegisteredResource>())
|
||||
public val tools: Map<String, RegisteredTool>
|
||||
get() = _tools.value
|
||||
public val prompts: Map<String, RegisteredPrompt>
|
||||
get() = _prompts.value
|
||||
public val resources: Map<String, RegisteredResource>
|
||||
get() = _resources.value
|
||||
|
||||
init {
|
||||
logger.debug { "Initializing MCP server with capabilities: $capabilities" }
|
||||
|
||||
// Core protocol handlers
|
||||
setRequestHandler<InitializeRequest>(Method.Defined.Initialize) { request, _ ->
|
||||
handleInitialize(request)
|
||||
}
|
||||
setNotificationHandler<InitializedNotification>(Method.Defined.NotificationsInitialized) {
|
||||
_onInitialized()
|
||||
CompletableDeferred(Unit)
|
||||
}
|
||||
|
||||
// Internal handlers for tools
|
||||
if (capabilities.tools != null) {
|
||||
setRequestHandler<ListToolsRequest>(Method.Defined.ToolsList) { _, _ ->
|
||||
handleListTools()
|
||||
}
|
||||
setRequestHandler<CallToolRequest>(Method.Defined.ToolsCall) { request, _ ->
|
||||
handleCallTool(request)
|
||||
}
|
||||
}
|
||||
|
||||
// Internal handlers for prompts
|
||||
if (capabilities.prompts != null) {
|
||||
setRequestHandler<ListPromptsRequest>(Method.Defined.PromptsList) { _, _ ->
|
||||
handleListPrompts()
|
||||
}
|
||||
setRequestHandler<GetPromptRequest>(Method.Defined.PromptsGet) { request, _ ->
|
||||
handleGetPrompt(request)
|
||||
}
|
||||
}
|
||||
|
||||
// Internal handlers for resources
|
||||
if (capabilities.resources != null) {
|
||||
setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
handleListResources()
|
||||
}
|
||||
setRequestHandler<ReadResourceRequest>(Method.Defined.ResourcesRead) { request, _ ->
|
||||
handleReadResource(request)
|
||||
}
|
||||
setRequestHandler<ListResourceTemplatesRequest>(Method.Defined.ResourcesTemplatesList) { _, _ ->
|
||||
handleListResourceTemplates()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a callback to be invoked when the server has completed initialization.
|
||||
*/
|
||||
public fun onInitialized(block: () -> Unit) {
|
||||
val old = _onInitialized
|
||||
_onInitialized = {
|
||||
old()
|
||||
block()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a callback to be invoked when the server connection is closing.
|
||||
*/
|
||||
public fun onClose(block: () -> Unit) {
|
||||
val old = _onClose
|
||||
_onClose = {
|
||||
old()
|
||||
block()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Called when the server connection is closing.
|
||||
*/
|
||||
override fun onClose() {
|
||||
logger.info { "Server connection closing" }
|
||||
_onClose()
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a single tool. The client can then call this tool.
|
||||
*
|
||||
* @param tool A [Tool] object describing the tool.
|
||||
* @param handler A suspend function that handles executing the tool when called by the client.
|
||||
* @throws IllegalStateException If the server does not support tools.
|
||||
*/
|
||||
public fun addTool(tool: Tool, handler: suspend (CallToolRequest) -> CallToolResult) {
|
||||
if (capabilities.tools == null) {
|
||||
logger.error { "Failed to add tool '${tool.name}': Server does not support tools capability" }
|
||||
throw IllegalStateException("Server does not support tools capability. Enable it in ServerOptions.")
|
||||
}
|
||||
logger.info { "Registering tool: ${tool.name}" }
|
||||
_tools.update { current -> current.put(tool.name, RegisteredTool(tool, handler)) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a single tool. The client can then call this tool.
|
||||
*
|
||||
* @param name The name of the tool.
|
||||
* @param title An optional human-readable name of the tool for display purposes.
|
||||
* @param description A human-readable description of what the tool does.
|
||||
* @param inputSchema The expected input schema for the tool.
|
||||
* @param outputSchema The optional expected output schema for the tool.
|
||||
* @param toolAnnotations Optional additional tool information.
|
||||
* @param handler A suspend function that handles executing the tool when called by the client.
|
||||
* @throws IllegalStateException If the server does not support tools.
|
||||
*/
|
||||
public fun addTool(
|
||||
name: String,
|
||||
description: String,
|
||||
inputSchema: Tool.Input = Tool.Input(),
|
||||
title: String? = null,
|
||||
outputSchema: Tool.Output? = null,
|
||||
toolAnnotations: ToolAnnotations? = null,
|
||||
handler: suspend (CallToolRequest) -> CallToolResult
|
||||
) {
|
||||
val tool = Tool(name, title, description, inputSchema, outputSchema, toolAnnotations)
|
||||
addTool(tool, handler)
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers multiple tools at once.
|
||||
*
|
||||
* @param toolsToAdd A list of [RegisteredTool] objects representing the tools to register.
|
||||
* @throws IllegalStateException If the server does not support tools.
|
||||
*/
|
||||
public fun addTools(toolsToAdd: List<RegisteredTool>) {
|
||||
if (capabilities.tools == null) {
|
||||
logger.error { "Failed to add tools: Server does not support tools capability" }
|
||||
throw IllegalStateException("Server does not support tools capability.")
|
||||
}
|
||||
logger.info { "Registering ${toolsToAdd.size} tools" }
|
||||
_tools.update { current -> current.putAll(toolsToAdd.associateBy { it.tool.name }) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes a single tool by name.
|
||||
*
|
||||
* @param name The name of the tool to remove.
|
||||
* @return True if the tool was removed, false if it wasn't found.
|
||||
* @throws IllegalStateException If the server does not support tools.
|
||||
*/
|
||||
public fun removeTool(name: String): Boolean {
|
||||
if (capabilities.tools == null) {
|
||||
logger.error { "Failed to remove tool '$name': Server does not support tools capability" }
|
||||
throw IllegalStateException("Server does not support tools capability.")
|
||||
}
|
||||
logger.info { "Removing tool: $name" }
|
||||
|
||||
val oldMap = _tools.getAndUpdate { current -> current.remove(name) }
|
||||
|
||||
val removed = name in oldMap
|
||||
logger.debug {
|
||||
if (removed) {
|
||||
"Tool removed: $name"
|
||||
} else {
|
||||
"Tool not found: $name"
|
||||
}
|
||||
}
|
||||
return removed
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes multiple tools at once.
|
||||
*
|
||||
* @param toolNames A list of tool names to remove.
|
||||
* @return The number of tools that were successfully removed.
|
||||
* @throws IllegalStateException If the server does not support tools.
|
||||
*/
|
||||
public fun removeTools(toolNames: List<String>): Int {
|
||||
if (capabilities.tools == null) {
|
||||
logger.error { "Failed to remove tools: Server does not support tools capability" }
|
||||
throw IllegalStateException("Server does not support tools capability.")
|
||||
}
|
||||
logger.info { "Removing ${toolNames.size} tools" }
|
||||
|
||||
val oldMap = _tools.getAndUpdate { current -> current - toolNames.toPersistentSet() }
|
||||
|
||||
val removedCount = toolNames.count { it in oldMap }
|
||||
logger.info {
|
||||
if (removedCount > 0) {
|
||||
"Removed $removedCount tools"
|
||||
} else {
|
||||
"No tools were removed"
|
||||
}
|
||||
}
|
||||
return removedCount
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a single prompt. The client can then retrieve the prompt.
|
||||
*
|
||||
* @param prompt A [Prompt] object describing the prompt.
|
||||
* @param promptProvider A suspend function that returns the prompt content when requested by the client.
|
||||
* @throws IllegalStateException If the server does not support prompts.
|
||||
*/
|
||||
public fun addPrompt(prompt: Prompt, promptProvider: suspend (GetPromptRequest) -> GetPromptResult) {
|
||||
if (capabilities.prompts == null) {
|
||||
logger.error { "Failed to add prompt '${prompt.name}': Server does not support prompts capability" }
|
||||
throw IllegalStateException("Server does not support prompts capability.")
|
||||
}
|
||||
logger.info { "Registering prompt: ${prompt.name}" }
|
||||
_prompts.update { current -> current.put(prompt.name, RegisteredPrompt(prompt, promptProvider)) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a single prompt by constructing a [Prompt] from given parameters.
|
||||
*
|
||||
* @param name The name of the prompt.
|
||||
* @param description An optional human-readable description of the prompt.
|
||||
* @param arguments An optional list of [PromptArgument] that the prompt accepts.
|
||||
* @param promptProvider A suspend function that returns the prompt content when requested.
|
||||
* @throws IllegalStateException If the server does not support prompts.
|
||||
*/
|
||||
public fun addPrompt(
|
||||
name: String,
|
||||
description: String? = null,
|
||||
arguments: List<PromptArgument>? = null,
|
||||
promptProvider: suspend (GetPromptRequest) -> GetPromptResult
|
||||
) {
|
||||
val prompt = Prompt(name = name, description = description, arguments = arguments)
|
||||
addPrompt(prompt, promptProvider)
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers multiple prompts at once.
|
||||
*
|
||||
* @param promptsToAdd A list of [RegisteredPrompt] objects representing the prompts to register.
|
||||
* @throws IllegalStateException If the server does not support prompts.
|
||||
*/
|
||||
public fun addPrompts(promptsToAdd: List<RegisteredPrompt>) {
|
||||
if (capabilities.prompts == null) {
|
||||
logger.error { "Failed to add prompts: Server does not support prompts capability" }
|
||||
throw IllegalStateException("Server does not support prompts capability.")
|
||||
}
|
||||
logger.info { "Registering ${promptsToAdd.size} prompts" }
|
||||
_prompts.update { current -> current.putAll(promptsToAdd.associateBy { it.prompt.name }) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes a single prompt by name.
|
||||
*
|
||||
* @param name The name of the prompt to remove.
|
||||
* @return True if the prompt was removed, false if it wasn't found.
|
||||
* @throws IllegalStateException If the server does not support prompts.
|
||||
*/
|
||||
public fun removePrompt(name: String): Boolean {
|
||||
if (capabilities.prompts == null) {
|
||||
logger.error { "Failed to remove prompt '$name': Server does not support prompts capability" }
|
||||
throw IllegalStateException("Server does not support prompts capability.")
|
||||
}
|
||||
logger.info { "Removing prompt: $name" }
|
||||
|
||||
val oldMap = _prompts.getAndUpdate { current -> current.remove(name) }
|
||||
|
||||
val removed = name in oldMap
|
||||
logger.debug {
|
||||
if (removed) {
|
||||
"Prompt removed: $name"
|
||||
} else {
|
||||
"Prompt not found: $name"
|
||||
}
|
||||
}
|
||||
return removed
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes multiple prompts at once.
|
||||
*
|
||||
* @param promptNames A list of prompt names to remove.
|
||||
* @return The number of prompts that were successfully removed.
|
||||
* @throws IllegalStateException If the server does not support prompts.
|
||||
*/
|
||||
public fun removePrompts(promptNames: List<String>): Int {
|
||||
if (capabilities.prompts == null) {
|
||||
logger.error { "Failed to remove prompts: Server does not support prompts capability" }
|
||||
throw IllegalStateException("Server does not support prompts capability.")
|
||||
}
|
||||
logger.info { "Removing ${promptNames.size} prompts" }
|
||||
|
||||
val oldMap = _prompts.getAndUpdate { current -> current - promptNames.toPersistentSet() }
|
||||
|
||||
val removedCount = promptNames.count { it in oldMap }
|
||||
|
||||
logger.info {
|
||||
if (removedCount > 0) {
|
||||
"Removed $removedCount prompts"
|
||||
} else {
|
||||
"No prompts were removed"
|
||||
}
|
||||
}
|
||||
return removedCount
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a single resource. The client can then read the resource content.
|
||||
*
|
||||
* @param uri The URI of the resource.
|
||||
* @param name A human-readable name for the resource.
|
||||
* @param description A description of the resource's content.
|
||||
* @param mimeType The MIME type of the resource content.
|
||||
* @param readHandler A suspend function that returns the resource content when read by the client.
|
||||
* @throws IllegalStateException If the server does not support resources.
|
||||
*/
|
||||
public fun addResource(
|
||||
uri: String,
|
||||
name: String,
|
||||
description: String,
|
||||
mimeType: String = "text/html",
|
||||
readHandler: suspend (ReadResourceRequest) -> ReadResourceResult
|
||||
) {
|
||||
if (capabilities.resources == null) {
|
||||
logger.error { "Failed to add resource '$name': Server does not support resources capability" }
|
||||
throw IllegalStateException("Server does not support resources capability.")
|
||||
}
|
||||
logger.info { "Registering resource: $name ($uri)" }
|
||||
_resources.update { current ->
|
||||
current.put(
|
||||
uri,
|
||||
RegisteredResource(Resource(uri, name, description, mimeType), readHandler)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers multiple resources at once.
|
||||
*
|
||||
* @param resourcesToAdd A list of [RegisteredResource] objects representing the resources to register.
|
||||
* @throws IllegalStateException If the server does not support resources.
|
||||
*/
|
||||
public fun addResources(resourcesToAdd: List<RegisteredResource>) {
|
||||
if (capabilities.resources == null) {
|
||||
logger.error { "Failed to add resources: Server does not support resources capability" }
|
||||
throw IllegalStateException("Server does not support resources capability.")
|
||||
}
|
||||
logger.info { "Registering ${resourcesToAdd.size} resources" }
|
||||
_resources.update { current -> current.putAll(resourcesToAdd.associateBy { it.resource.uri }) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes a single resource by URI.
|
||||
*
|
||||
* @param uri The URI of the resource to remove.
|
||||
* @return True if the resource was removed, false if it wasn't found.
|
||||
* @throws IllegalStateException If the server does not support resources.
|
||||
*/
|
||||
public fun removeResource(uri: String): Boolean {
|
||||
if (capabilities.resources == null) {
|
||||
logger.error { "Failed to remove resource '$uri': Server does not support resources capability" }
|
||||
throw IllegalStateException("Server does not support resources capability.")
|
||||
}
|
||||
logger.info { "Removing resource: $uri" }
|
||||
|
||||
val oldMap = _resources.getAndUpdate { current -> current.remove(uri) }
|
||||
|
||||
val removed = uri in oldMap
|
||||
logger.debug {
|
||||
if (removed) {
|
||||
"Resource removed: $uri"
|
||||
} else {
|
||||
"Resource not found: $uri"
|
||||
}
|
||||
}
|
||||
return removed
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes multiple resources at once.
|
||||
*
|
||||
* @param uris A list of resource URIs to remove.
|
||||
* @return The number of resources that were successfully removed.
|
||||
* @throws IllegalStateException If the server does not support resources.
|
||||
*/
|
||||
public fun removeResources(uris: List<String>): Int {
|
||||
if (capabilities.resources == null) {
|
||||
logger.error { "Failed to remove resources: Server does not support resources capability" }
|
||||
throw IllegalStateException("Server does not support resources capability.")
|
||||
}
|
||||
logger.info { "Removing ${uris.size} resources" }
|
||||
|
||||
val oldMap = _resources.getAndUpdate { current -> current - uris.toPersistentSet() }
|
||||
|
||||
val removedCount = uris.count { it in oldMap }
|
||||
|
||||
logger.info {
|
||||
if (removedCount > 0) {
|
||||
"Removed $removedCount resources"
|
||||
} else {
|
||||
"No resources were removed"
|
||||
}
|
||||
}
|
||||
return removedCount
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a ping request to the client to check connectivity.
|
||||
*
|
||||
* @return The result of the ping request.
|
||||
* @throws IllegalStateException If for some reason the method is not supported or the connection is closed.
|
||||
*/
|
||||
public suspend fun ping(): EmptyRequestResult {
|
||||
return request<EmptyRequestResult>(PingRequest())
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a message using the server's sampling capability.
|
||||
*
|
||||
* @param params The parameters for creating a message.
|
||||
* @param options Optional request options.
|
||||
* @return The created message result.
|
||||
* @throws IllegalStateException If the server does not support sampling or if the request fails.
|
||||
*/
|
||||
public suspend fun createMessage(
|
||||
params: CreateMessageRequest,
|
||||
options: RequestOptions? = null
|
||||
): CreateMessageResult {
|
||||
logger.debug { "Creating message with params: $params" }
|
||||
return request<CreateMessageResult>(params, options)
|
||||
}
|
||||
|
||||
/**
|
||||
* Lists the available "roots" from the client's perspective (if supported).
|
||||
*
|
||||
* @param params JSON parameters for the request, usually empty.
|
||||
* @param options Optional request options.
|
||||
* @return The list of roots.
|
||||
* @throws IllegalStateException If the server or client does not support roots.
|
||||
*/
|
||||
public suspend fun listRoots(
|
||||
params: JsonObject = EmptyJsonObject,
|
||||
options: RequestOptions? = null
|
||||
): ListRootsResult {
|
||||
logger.debug { "Listing roots with params: $params" }
|
||||
return request<ListRootsResult>(ListRootsRequest(params), options)
|
||||
}
|
||||
|
||||
public suspend fun createElicitation(
|
||||
message: String,
|
||||
requestedSchema: RequestedSchema,
|
||||
options: RequestOptions? = null
|
||||
): CreateElicitationResult {
|
||||
logger.debug { "Creating elicitation with message: $message" }
|
||||
return request(CreateElicitationRequest(message, requestedSchema), options)
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a logging message notification to the client.
|
||||
*
|
||||
* @param params The logging message notification parameters.
|
||||
*/
|
||||
public suspend fun sendLoggingMessage(params: LoggingMessageNotification) {
|
||||
logger.trace { "Sending logging message: ${params.params.data}" }
|
||||
notification(params)
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a resource-updated notification to the client, indicating that a specific resource has changed.
|
||||
*
|
||||
* @param params Details of the updated resource.
|
||||
*/
|
||||
public suspend fun sendResourceUpdated(params: ResourceUpdatedNotification) {
|
||||
logger.debug { "Sending resource updated notification for: ${params.params.uri}" }
|
||||
notification(params)
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a notification to the client indicating that the list of resources has changed.
|
||||
*/
|
||||
public suspend fun sendResourceListChanged() {
|
||||
logger.debug { "Sending resource list changed notification" }
|
||||
notification(ResourceListChangedNotification())
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a notification to the client indicating that the list of tools has changed.
|
||||
*/
|
||||
public suspend fun sendToolListChanged() {
|
||||
logger.debug { "Sending tool list changed notification" }
|
||||
notification(ToolListChangedNotification())
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a notification to the client indicating that the list of prompts has changed.
|
||||
*/
|
||||
public suspend fun sendPromptListChanged() {
|
||||
logger.debug { "Sending prompt list changed notification" }
|
||||
notification(PromptListChangedNotification())
|
||||
}
|
||||
|
||||
// --- Internal Handlers ---
|
||||
|
||||
private suspend fun handleInitialize(request: InitializeRequest): InitializeResult {
|
||||
logger.info { "Handling initialize request from client ${request.clientInfo}" }
|
||||
clientCapabilities = request.capabilities
|
||||
clientVersion = request.clientInfo
|
||||
|
||||
val requestedVersion = request.protocolVersion
|
||||
val protocolVersion = if (SUPPORTED_PROTOCOL_VERSIONS.contains(requestedVersion)) {
|
||||
requestedVersion
|
||||
} else {
|
||||
logger.warn { "Client requested unsupported protocol version $requestedVersion, falling back to $LATEST_PROTOCOL_VERSION" }
|
||||
LATEST_PROTOCOL_VERSION
|
||||
}
|
||||
|
||||
return InitializeResult(
|
||||
protocolVersion = protocolVersion,
|
||||
capabilities = capabilities,
|
||||
serverInfo = serverInfo
|
||||
)
|
||||
}
|
||||
|
||||
private suspend fun handleListTools(): ListToolsResult {
|
||||
val toolList = tools.values.map { it.tool }
|
||||
return ListToolsResult(tools = toolList, nextCursor = null)
|
||||
}
|
||||
|
||||
private suspend fun handleCallTool(request: CallToolRequest): CallToolResult {
|
||||
logger.debug { "Handling tool call request for tool: ${request.name}" }
|
||||
val tool = _tools.value[request.name]
|
||||
?: run {
|
||||
logger.error { "Tool not found: ${request.name}" }
|
||||
throw IllegalArgumentException("Tool not found: ${request.name}")
|
||||
}
|
||||
logger.trace { "Executing tool ${request.name} with input: ${request.arguments}" }
|
||||
return tool.handler(request)
|
||||
}
|
||||
|
||||
private suspend fun handleListPrompts(): ListPromptsResult {
|
||||
logger.debug { "Handling list prompts request" }
|
||||
return ListPromptsResult(prompts = prompts.values.map { it.prompt })
|
||||
}
|
||||
|
||||
private suspend fun handleGetPrompt(request: GetPromptRequest): GetPromptResult {
|
||||
logger.debug { "Handling get prompt request for: ${request.name}" }
|
||||
val prompt = prompts[request.name]
|
||||
?: run {
|
||||
logger.error { "Prompt not found: ${request.name}" }
|
||||
throw IllegalArgumentException("Prompt not found: ${request.name}")
|
||||
}
|
||||
return prompt.messageProvider(request)
|
||||
}
|
||||
|
||||
private suspend fun handleListResources(): ListResourcesResult {
|
||||
logger.debug { "Handling list resources request" }
|
||||
return ListResourcesResult(resources = resources.values.map { it.resource })
|
||||
}
|
||||
|
||||
private suspend fun handleReadResource(request: ReadResourceRequest): ReadResourceResult {
|
||||
logger.debug { "Handling read resource request for: ${request.uri}" }
|
||||
val resource = resources[request.uri]
|
||||
?: run {
|
||||
logger.error { "Resource not found: ${request.uri}" }
|
||||
throw IllegalArgumentException("Resource not found: ${request.uri}")
|
||||
}
|
||||
return resource.readHandler(request)
|
||||
}
|
||||
|
||||
private suspend fun handleListResourceTemplates(): ListResourceTemplatesResult {
|
||||
// If you have resource templates, return them here. For now, return empty.
|
||||
return ListResourceTemplatesResult(listOf())
|
||||
}
|
||||
|
||||
/**
|
||||
* Asserts that the client supports the capability required for the given [method].
|
||||
*
|
||||
* This method is automatically called by the [Protocol] framework before handling requests.
|
||||
* Throws [IllegalStateException] if the capability is not supported.
|
||||
*
|
||||
* @param method The method for which we are asserting capability.
|
||||
*/
|
||||
override fun assertCapabilityForMethod(method: Method) {
|
||||
logger.trace { "Asserting capability for method: ${method.value}" }
|
||||
when (method.value) {
|
||||
"sampling/createMessage" -> {
|
||||
if (clientCapabilities?.sampling == null) {
|
||||
logger.error { "Client capability assertion failed: sampling not supported" }
|
||||
throw IllegalStateException("Client does not support sampling (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"roots/list" -> {
|
||||
if (clientCapabilities?.roots == null) {
|
||||
throw IllegalStateException("Client does not support listing roots (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"elicitation/create" -> {
|
||||
if (clientCapabilities?.elicitation == null) {
|
||||
throw IllegalStateException("Client does not support elicitation (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"ping" -> {
|
||||
// No specific capability required
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Asserts that the server can handle the specified notification method.
|
||||
*
|
||||
* Throws [IllegalStateException] if the server does not have the capabilities required to handle this notification.
|
||||
*
|
||||
* @param method The notification method.
|
||||
*/
|
||||
override fun assertNotificationCapability(method: Method) {
|
||||
logger.trace { "Asserting notification capability for method: ${method.value}" }
|
||||
when (method.value) {
|
||||
"notifications/message" -> {
|
||||
if (capabilities.logging == null) {
|
||||
logger.error { "Server capability assertion failed: logging not supported" }
|
||||
throw IllegalStateException("Server does not support logging (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"notifications/resources/updated",
|
||||
"notifications/resources/list_changed" -> {
|
||||
if (capabilities.resources == null) {
|
||||
throw IllegalStateException("Server does not support notifying about resources (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"notifications/tools/list_changed" -> {
|
||||
if (capabilities.tools == null) {
|
||||
throw IllegalStateException("Server does not support notifying of tool list changes (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"notifications/prompts/list_changed" -> {
|
||||
if (capabilities.prompts == null) {
|
||||
throw IllegalStateException("Server does not support notifying of prompt list changes (required for ${method.value})")
|
||||
}
|
||||
}
|
||||
|
||||
"notifications/cancelled",
|
||||
"notifications/progress" -> {
|
||||
// Always allowed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Asserts that the server can handle the specified request method.
|
||||
*
|
||||
* Throws [IllegalStateException] if the server does not have the capabilities required to handle this request.
|
||||
*
|
||||
* @param method The request method.
|
||||
*/
|
||||
override fun assertRequestHandlerCapability(method: Method) {
|
||||
logger.trace { "Asserting request handler capability for method: ${method.value}" }
|
||||
when (method.value) {
|
||||
"sampling/createMessage" -> {
|
||||
if (capabilities.sampling == null) {
|
||||
logger.error { "Server capability assertion failed: sampling not supported" }
|
||||
throw IllegalStateException("Server does not support sampling (required for $method)")
|
||||
}
|
||||
}
|
||||
|
||||
"logging/setLevel" -> {
|
||||
if (capabilities.logging == null) {
|
||||
throw IllegalStateException("Server does not support logging (required for $method)")
|
||||
}
|
||||
}
|
||||
|
||||
"prompts/get",
|
||||
"prompts/list" -> {
|
||||
if (capabilities.prompts == null) {
|
||||
throw IllegalStateException("Server does not support prompts (required for $method)")
|
||||
}
|
||||
}
|
||||
|
||||
"resources/list",
|
||||
"resources/templates/list",
|
||||
"resources/read" -> {
|
||||
if (capabilities.resources == null) {
|
||||
throw IllegalStateException("Server does not support resources (required for $method)")
|
||||
}
|
||||
}
|
||||
|
||||
"tools/call",
|
||||
"tools/list" -> {
|
||||
if (capabilities.tools == null) {
|
||||
throw IllegalStateException("Server does not support tools (required for $method)")
|
||||
}
|
||||
}
|
||||
|
||||
"ping", "initialize" -> {
|
||||
// No capability required
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A wrapper class representing a registered tool on the server.
|
||||
*
|
||||
* @property tool The tool definition.
|
||||
* @property handler A suspend function to handle the tool call requests.
|
||||
*/
|
||||
public data class RegisteredTool(
|
||||
val tool: Tool,
|
||||
val handler: suspend (CallToolRequest) -> CallToolResult
|
||||
)
|
||||
|
||||
/**
|
||||
* A wrapper class representing a registered prompt on the server.
|
||||
*
|
||||
* @property prompt The prompt definition.
|
||||
* @property messageProvider A suspend function that returns the prompt content when requested by the client.
|
||||
*/
|
||||
public data class RegisteredPrompt(
|
||||
val prompt: Prompt,
|
||||
val messageProvider: suspend (GetPromptRequest) -> GetPromptResult
|
||||
)
|
||||
|
||||
/**
|
||||
* A wrapper class representing a registered resource on the server.
|
||||
*
|
||||
* @property resource The resource definition.
|
||||
* @property readHandler A suspend function to handle read requests for this resource.
|
||||
*/
|
||||
public data class RegisteredResource(
|
||||
val resource: Resource,
|
||||
val readHandler: suspend (ReadResourceRequest) -> ReadResourceResult
|
||||
)
|
||||
|
|
@ -0,0 +1,154 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.internal.IODispatcher
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
|
||||
import kotlinx.coroutines.CoroutineScope
|
||||
import kotlinx.coroutines.Job
|
||||
import kotlinx.coroutines.NonCancellable
|
||||
import kotlinx.coroutines.SupervisorJob
|
||||
import kotlinx.coroutines.cancelAndJoin
|
||||
import kotlinx.coroutines.channels.Channel
|
||||
import kotlinx.coroutines.isActive
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.withContext
|
||||
import kotlinx.io.Buffer
|
||||
import kotlinx.io.Sink
|
||||
import kotlinx.io.Source
|
||||
import kotlinx.io.buffered
|
||||
import kotlinx.io.readByteArray
|
||||
import kotlinx.io.writeString
|
||||
import kotlin.concurrent.atomics.AtomicBoolean
|
||||
import kotlin.concurrent.atomics.ExperimentalAtomicApi
|
||||
import kotlin.coroutines.CoroutineContext
|
||||
|
||||
/**
|
||||
* A server transport that communicates with a client via standard I/O.
|
||||
*
|
||||
* Reads from System.in and writes to System.out.
|
||||
*/
|
||||
@OptIn(ExperimentalAtomicApi::class)
|
||||
public class StdioServerTransport(
|
||||
private val inputStream: Source,
|
||||
outputStream: Sink
|
||||
) : AbstractTransport() {
|
||||
private val logger = KotlinLogging.logger {}
|
||||
|
||||
private val readBuffer = ReadBuffer()
|
||||
private val initialized: AtomicBoolean = AtomicBoolean(false)
|
||||
private var readingJob: Job? = null
|
||||
private var sendingJob: Job? = null
|
||||
|
||||
private val coroutineContext: CoroutineContext = IODispatcher + SupervisorJob()
|
||||
private val scope = CoroutineScope(coroutineContext)
|
||||
private val readChannel = Channel<ByteArray>(Channel.UNLIMITED)
|
||||
private val writeChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
private val outputWriter = outputStream.buffered()
|
||||
|
||||
override suspend fun start() {
|
||||
if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
|
||||
error("StdioServerTransport already started!")
|
||||
}
|
||||
|
||||
// Launch a coroutine to read from stdin
|
||||
readingJob = scope.launch {
|
||||
val buf = Buffer()
|
||||
try {
|
||||
while (isActive) {
|
||||
val bytesRead = inputStream.readAtMostTo(buf, 8192)
|
||||
if (bytesRead == -1L) {
|
||||
// EOF reached
|
||||
break
|
||||
}
|
||||
if (bytesRead > 0) {
|
||||
val chunk = buf.readByteArray()
|
||||
readChannel.send(chunk)
|
||||
}
|
||||
}
|
||||
} catch (e: Throwable) {
|
||||
logger.error(e) { "Error reading from stdin" }
|
||||
_onError.invoke(e)
|
||||
} finally {
|
||||
// Reached EOF or error, close connection
|
||||
close()
|
||||
}
|
||||
}
|
||||
|
||||
// Launch a coroutine to process messages from readChannel
|
||||
scope.launch {
|
||||
try {
|
||||
for (chunk in readChannel) {
|
||||
readBuffer.append(chunk)
|
||||
processReadBuffer()
|
||||
}
|
||||
} catch (e: Throwable) {
|
||||
_onError.invoke(e)
|
||||
}
|
||||
}
|
||||
|
||||
// Launch a coroutine to handle message sending
|
||||
sendingJob = scope.launch {
|
||||
try {
|
||||
for (message in writeChannel) {
|
||||
val json = serializeMessage(message)
|
||||
outputWriter.writeString(json)
|
||||
outputWriter.flush()
|
||||
}
|
||||
} catch (e: Throwable) {
|
||||
logger.error(e) { "Error writing to stdout" }
|
||||
_onError.invoke(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun processReadBuffer() {
|
||||
while (true) {
|
||||
val message = try {
|
||||
readBuffer.readMessage()
|
||||
} catch (e: Throwable) {
|
||||
_onError.invoke(e)
|
||||
null
|
||||
}
|
||||
|
||||
if (message == null) break
|
||||
// Async invocation broke delivery order
|
||||
try {
|
||||
_onMessage.invoke(message)
|
||||
} catch (e: Throwable) {
|
||||
_onError.invoke(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
if (!initialized.compareAndSet(expectedValue = true, newValue = false)) return
|
||||
|
||||
withContext(NonCancellable) {
|
||||
writeChannel.close()
|
||||
sendingJob?.cancelAndJoin()
|
||||
|
||||
runCatching {
|
||||
inputStream.close()
|
||||
}.onFailure { logger.warn(it) { "Failed to close stdin" } }
|
||||
|
||||
readingJob?.cancel()
|
||||
|
||||
readChannel.close()
|
||||
readBuffer.clear()
|
||||
|
||||
runCatching {
|
||||
outputWriter.flush()
|
||||
outputWriter.close()
|
||||
}.onFailure { logger.warn(it) { "Failed to close stdout" } }
|
||||
|
||||
_onClose.invoke()
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage) {
|
||||
writeChannel.send(message)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.ktor.server.routing.Route
|
||||
import io.ktor.server.websocket.WebSocketServerSession
|
||||
import io.ktor.server.websocket.webSocket
|
||||
import io.modelcontextprotocol.kotlin.sdk.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.LIB_VERSION
|
||||
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.IMPLEMENTATION_NAME
|
||||
|
||||
/**
|
||||
* Registers a WebSocket route that establishes an MCP (Model Context Protocol) server session.
|
||||
*
|
||||
* @param options Optional server configuration settings for the MCP server.
|
||||
* @param handler A suspend function that defines the server's behavior.
|
||||
*/
|
||||
public fun Route.mcpWebSocket(
|
||||
options: ServerOptions? = null,
|
||||
handler: suspend Server.() -> Unit = {},
|
||||
) {
|
||||
webSocket {
|
||||
createMcpServer(this, options, handler)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a WebSocket route at the specified [path] that establishes an MCP server session.
|
||||
*
|
||||
* @param path The URL path at which to register the WebSocket route.
|
||||
* @param options Optional server configuration settings for the MCP server.
|
||||
* @param handler A suspend function that defines the server's behavior.
|
||||
*/
|
||||
public fun Route.mcpWebSocket(
|
||||
path: String,
|
||||
options: ServerOptions? = null,
|
||||
handler: suspend Server.() -> Unit = {},
|
||||
) {
|
||||
webSocket(path) {
|
||||
createMcpServer(this, options, handler)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a WebSocket route that creates an MCP server transport layer.
|
||||
*
|
||||
* @param handler A suspend function that defines the behavior of the transport layer.
|
||||
*/
|
||||
public fun Route.mcpWebSocketTransport(
|
||||
handler: suspend WebSocketMcpServerTransport.() -> Unit = {},
|
||||
) {
|
||||
webSocket {
|
||||
val transport = createMcpTransport(this)
|
||||
transport.start()
|
||||
handler(transport)
|
||||
transport.close()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a WebSocket route at the specified [path] that creates an MCP server transport layer.
|
||||
*
|
||||
* @param path The URL path at which to register the WebSocket route.
|
||||
* @param handler A suspend function that defines the behavior of the transport layer.
|
||||
*/
|
||||
public fun Route.mcpWebSocketTransport(
|
||||
path: String,
|
||||
handler: suspend WebSocketMcpServerTransport.() -> Unit = {},
|
||||
) {
|
||||
webSocket(path) {
|
||||
val transport = createMcpTransport(this)
|
||||
transport.start()
|
||||
handler(transport)
|
||||
transport.close()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
private suspend fun Route.createMcpServer(
|
||||
session: WebSocketServerSession,
|
||||
options: ServerOptions?,
|
||||
handler: suspend Server.() -> Unit,
|
||||
) {
|
||||
val transport = createMcpTransport(session)
|
||||
|
||||
val server = Server(
|
||||
serverInfo = Implementation(
|
||||
name = IMPLEMENTATION_NAME,
|
||||
version = LIB_VERSION
|
||||
),
|
||||
options = options ?: ServerOptions(
|
||||
capabilities = ServerCapabilities(
|
||||
prompts = ServerCapabilities.Prompts(listChanged = null),
|
||||
resources = ServerCapabilities.Resources(subscribe = null, listChanged = null),
|
||||
tools = ServerCapabilities.Tools(listChanged = null),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
server.connect(transport)
|
||||
handler(server)
|
||||
server.close()
|
||||
}
|
||||
|
||||
private fun createMcpTransport(
|
||||
session: WebSocketServerSession,
|
||||
): WebSocketMcpServerTransport {
|
||||
return WebSocketMcpServerTransport(session)
|
||||
}
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.ktor.http.HttpHeaders
|
||||
import io.ktor.server.websocket.WebSocketServerSession
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.MCP_SUBPROTOCOL
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.WebSocketMcpTransport
|
||||
|
||||
/**
|
||||
* Server-side implementation of the MCP (Model Context Protocol) transport over WebSocket.
|
||||
*
|
||||
* @property session The WebSocket server session used for communication.
|
||||
*/
|
||||
public class WebSocketMcpServerTransport(
|
||||
override val session: WebSocketServerSession,
|
||||
) : WebSocketMcpTransport() {
|
||||
override suspend fun initializeSession() {
|
||||
val subprotocol = session.call.request.headers[HttpHeaders.SecWebSocketProtocol]
|
||||
if (subprotocol != MCP_SUBPROTOCOL) {
|
||||
error("Invalid subprotocol: $subprotocol, expected $MCP_SUBPROTOCOL")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,144 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.InitializedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.PingRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.toJSON
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.io.Sink
|
||||
import kotlinx.io.Source
|
||||
import kotlinx.io.asSink
|
||||
import kotlinx.io.asSource
|
||||
import kotlinx.io.buffered
|
||||
import java.io.ByteArrayOutputStream
|
||||
import java.io.PipedInputStream
|
||||
import java.io.PipedOutputStream
|
||||
import kotlin.test.BeforeTest
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertFalse
|
||||
import kotlin.test.assertTrue
|
||||
|
||||
class StdioServerTransportTest {
|
||||
private lateinit var input: PipedInputStream
|
||||
private lateinit var inputWriter: PipedOutputStream
|
||||
private lateinit var outputBuffer: ReadBuffer
|
||||
private lateinit var output: ByteArrayOutputStream
|
||||
|
||||
// We'll store the wrapped streams that meet the constructor requirements
|
||||
private lateinit var bufferedInput: Source
|
||||
private lateinit var printOutput: Sink
|
||||
|
||||
@BeforeTest
|
||||
fun setUp() {
|
||||
// Simulate an input stream that we can push data into using inputWriter.
|
||||
input = PipedInputStream()
|
||||
inputWriter = PipedOutputStream(input)
|
||||
|
||||
outputBuffer = ReadBuffer()
|
||||
|
||||
// A custom ByteArrayOutputStream that appends all written data into outputBuffer
|
||||
output = object : ByteArrayOutputStream() {
|
||||
override fun write(b: ByteArray, off: Int, len: Int) {
|
||||
super.write(b, off, len)
|
||||
outputBuffer.append(b.copyOfRange(off, off + len))
|
||||
}
|
||||
}
|
||||
|
||||
bufferedInput = input.asSource().buffered()
|
||||
|
||||
printOutput = output.asSink().buffered()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should start then close cleanly`() {
|
||||
runBlocking {
|
||||
val server = StdioServerTransport(bufferedInput, printOutput)
|
||||
server.onError { error ->
|
||||
throw error
|
||||
}
|
||||
|
||||
var didClose = false
|
||||
server.onClose {
|
||||
didClose = true
|
||||
}
|
||||
|
||||
server.start()
|
||||
assertFalse(didClose, "Should not have closed yet")
|
||||
|
||||
server.close()
|
||||
assertTrue(didClose, "Should have closed after calling close()")
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should not read until started`() = runTest {
|
||||
val server = StdioServerTransport(bufferedInput, printOutput)
|
||||
server.onError { error ->
|
||||
throw error
|
||||
}
|
||||
|
||||
var didRead = false
|
||||
val readMessage = CompletableDeferred<JSONRPCMessage>()
|
||||
|
||||
server.onMessage { message ->
|
||||
didRead = true
|
||||
readMessage.complete(message)
|
||||
}
|
||||
|
||||
val message = PingRequest().toJSON()
|
||||
|
||||
// Push a message before the server started
|
||||
val serialized = serializeMessage(message)
|
||||
inputWriter.write(serialized)
|
||||
inputWriter.flush()
|
||||
|
||||
assertFalse(didRead, "Should not have read message before start")
|
||||
|
||||
server.start()
|
||||
val received = readMessage.await()
|
||||
assertEquals(message, received)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `should read multiple messages`() = runTest {
|
||||
val server = StdioServerTransport(bufferedInput, printOutput)
|
||||
server.onError { error ->
|
||||
throw error
|
||||
}
|
||||
|
||||
val messages = listOf(
|
||||
PingRequest().toJSON(),
|
||||
InitializedNotification().toJSON(),
|
||||
)
|
||||
|
||||
val readMessages = mutableListOf<JSONRPCMessage>()
|
||||
val finished = CompletableDeferred<Unit>()
|
||||
|
||||
server.onMessage { message ->
|
||||
readMessages.add(message)
|
||||
if (message == messages[1]) {
|
||||
finished.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
// Push both messages before starting the server
|
||||
for (m in messages) {
|
||||
inputWriter.write(serializeMessage(m))
|
||||
}
|
||||
inputWriter.flush()
|
||||
|
||||
server.start()
|
||||
finished.await()
|
||||
|
||||
assertEquals(messages, readMessages)
|
||||
}
|
||||
}
|
||||
|
||||
fun PipedOutputStream.write(s: String) {
|
||||
write(s.toByteArray())
|
||||
}
|
||||
Loading…
Reference in a new issue