modelcontextprotocol/kotlin-sdk
Watch
1
0
Fork
You've already forked kotlin-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Initial commit

Browse source
This commit is contained in:
leonid.stashevsky 2024-12-15 20:17:55 +01:00
commit 7b32c3e3da
30 changed files with 7267 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

419
src/main/kotlin/org/jetbrains/kotlinx/mcp/client/Client.kt Normal file
View file

@ -0,0 +1,419 @@
package org.jetbrains.kotlinx.mcp.client
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import org.jetbrains.kotlinx.mcp.*
import org.jetbrains.kotlinx.mcp.LoggingMessageNotification.SetLevelRequest
import org.jetbrains.kotlinx.mcp.shared.Protocol
import org.jetbrains.kotlinx.mcp.shared.ProtocolOptions
import org.jetbrains.kotlinx.mcp.shared.RequestOptions
import org.jetbrains.kotlinx.mcp.shared.Transport
/**
* Options for configuring the MCP client.
*
* @property capabilities The capabilities this client supports.
* @property enforceStrictCapabilities Whether to strictly enforce capabilities when interacting with the server.
*/
class ClientOptions(
val capabilities: ClientCapabilities = ClientCapabilities(),
enforceStrictCapabilities: Boolean = true,
) : ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities)
/**
* An MCP client on top of a pluggable transport.
*
* The client automatically performs the initialization handshake with the server when [connect] is called.
* After initialization, [getServerCapabilities] and [getServerVersion] provide details about the connected server.
*
* You can extend this class with custom request/notification/result types if needed.
*
* @param clientInfo Information about the client implementation (name, version).
* @param options Configuration options for this client.
*/
open class Client(
private val clientInfo: Implementation,
options: ClientOptions = ClientOptions(),
) : Protocol<ClientRequest, ClientNotification, ClientResult>(options) {
private var serverCapabilities: ServerCapabilities? = null
private var serverVersion: Implementation? = null
private val capabilities: ClientCapabilities = options.capabilities
protected fun assertCapability(capability: String, method: String) {
val caps = serverCapabilities
val hasCapability = when (capability) {
"logging" -> caps?.logging != null
"prompts" -> caps?.prompts != null
"resources" -> caps?.resources != null
"tools" -> caps?.tools != null
else -> true
}
if (!hasCapability) {
throw IllegalStateException("Server does not support $capability (required for $method)")
}
}
/**
* Connects the client to the given [transport], performing the initialization handshake with the server.
*
* @param transport The transport to use for communication with the server.
* @throws IllegalStateException If the server's protocol version is not supported.
*/
override suspend fun connect(transport: Transport) {
super.connect(transport)
try {
val message = InitializeRequest(
protocolVersion = LATEST_PROTOCOL_VERSION,
capabilities = capabilities,
clientInfo = clientInfo
)
val result = request<InitializeResult>(message)
if (!SUPPORTED_PROTOCOL_VERSIONS.contains(result.protocolVersion)) {
throw IllegalStateException(
"Server's protocol version is not supported: ${result.protocolVersion}"
)
}
serverCapabilities = result.capabilities
serverVersion = result.serverInfo
notification(InitializedNotification())
} catch (error: Throwable) {
close()
throw error
}
}
/**
* Retrieves the server's reported capabilities after the initialization process completes.
*
* @return The server's capabilities, or `null` if initialization is not yet complete.
*/
fun getServerCapabilities(): ServerCapabilities? {
return serverCapabilities
}
/**
* Retrieves the server's reported version information after initialization.
*
* @return Information about the server's implementation, or `null` if initialization is not yet complete.
*/
fun getServerVersion(): Implementation? {
return serverVersion
}
override fun assertCapabilityForMethod(method: Method) {
when (method) {
Method.Defined.LoggingSetLevel -> {
if (serverCapabilities?.logging == null) {
throw IllegalStateException("Server does not support logging (required for $method)")
}
}
Method.Defined.PromptsGet,
Method.Defined.PromptsList,
Method.Defined.CompletionComplete,
-> {
if (serverCapabilities?.prompts == null) {
throw IllegalStateException("Server does not support prompts (required for $method)")
}
}
Method.Defined.ResourcesList,
Method.Defined.ResourcesTemplatesList,
Method.Defined.ResourcesRead,
Method.Defined.ResourcesSubscribe,
Method.Defined.ResourcesUnsubscribe,
-> {
val resCaps = serverCapabilities?.resources
if (resCaps == null) {
throw IllegalStateException("Server does not support resources (required for $method)")
}
if (method == Method.Defined.ResourcesSubscribe && resCaps.subscribe != true) {
throw IllegalStateException(
"Server does not support resource subscriptions (required for $method)"
)
}
}
Method.Defined.ToolsCall,
Method.Defined.ToolsList,
-> {
if (serverCapabilities?.tools == null) {
throw IllegalStateException("Server does not support tools (required for $method)")
}
}
Method.Defined.Initialize,
Method.Defined.Ping,
-> {
// No specific capability required
}
else -> {
// For unknown or future methods, no assertion by default
}
}
}
override fun assertNotificationCapability(method: Method) {
when (method) {
Method.Defined.NotificationsRootsListChanged -> {
if (capabilities.roots?.listChanged != true) {
throw IllegalStateException(
"Client does not support roots list changed notifications (required for $method)"
)
}
}
Method.Defined.NotificationsInitialized,
Method.Defined.NotificationsCancelled,
Method.Defined.NotificationsProgress,
-> {
// Always allowed
}
else -> {
// For notifications not specifically listed, no assertion by default
}
}
}
override fun assertRequestHandlerCapability(method: Method) {
when (method) {
Method.Defined.SamplingCreateMessage -> {
if (capabilities.sampling == null) {
throw IllegalStateException(
"Client does not support sampling capability (required for $method)"
)
}
}
Method.Defined.RootsList -> {
if (capabilities.roots == null) {
throw IllegalStateException(
"Client does not support roots capability (required for $method)"
)
}
}
Method.Defined.Ping -> {
// No capability required
}
else -> {}
}
}
/**
* Sends a ping request to the server to check connectivity.
*
* @param options Optional request options.
* @throws IllegalStateException If the server does not support the ping method (unlikely).
*/
suspend fun ping(options: RequestOptions? = null): EmptyRequestResult {
return request<EmptyRequestResult>(PingRequest(), options)
}
/**
* Sends a completion request to the server, typically to generate or complete some content.
*
* @param params The completion request parameters.
* @param options Optional request options.
* @return The completion result returned by the server, or `null` if none.
* @throws IllegalStateException If the server does not support prompts or completion.
*/
suspend fun complete(params: CompleteRequest, options: RequestOptions? = null): CompleteResult? {
return request<CompleteResult>(params, options)
}
/**
* Sets the logging level on the server.
*
* @param level The desired logging level.
* @param options Optional request options.
* @throws IllegalStateException If the server does not support logging.
*/
suspend fun setLoggingLevel(level: LoggingLevel, options: RequestOptions? = null): EmptyRequestResult {
return request<EmptyRequestResult>(SetLevelRequest(level), options)
}
/**
* Retrieves a prompt by name from the server.
*
* @param request The prompt request containing the prompt name.
* @param options Optional request options.
* @return The requested prompt details, or `null` if not found.
* @throws IllegalStateException If the server does not support prompts.
*/
suspend fun getPrompt(request: GetPromptRequest, options: RequestOptions? = null): GetPromptResult? {
return request<GetPromptResult>(request, options)
}
/**
* Lists all available prompts from the server.
*
* @param request A request object for listing prompts (usually empty).
* @param options Optional request options.
* @return The list of available prompts, or `null` if none.
* @throws IllegalStateException If the server does not support prompts.
*/
suspend fun listPrompts(
request: ListPromptsRequest = ListPromptsRequest(),
options: RequestOptions? = null,
): ListPromptsResult? {
return request<ListPromptsResult>(request, options)
}
/**
* Lists all available resources from the server.
*
* @param request A request object for listing resources (usually empty).
* @param options Optional request options.
* @return The list of resources, or `null` if none.
* @throws IllegalStateException If the server does not support resources.
*/
suspend fun listResources(
request: ListResourcesRequest = ListResourcesRequest(),
options: RequestOptions? = null,
): ListResourcesResult? {
return request<ListResourcesResult>(request, options)
}
/**
* Lists resource templates available on the server.
*
* @param request The request object for listing resource templates.
* @param options Optional request options.
* @return The list of resource templates, or `null` if none.
* @throws IllegalStateException If the server does not support resources.
*/
suspend fun listResourceTemplates(
request: ListResourceTemplatesRequest,
options: RequestOptions? = null,
): ListResourceTemplatesResult? {
return request<ListResourceTemplatesResult>(request, options)
}
/**
* Reads a resource from the server by its URI.
*
* @param request The request object containing the resource URI.
* @param options Optional request options.
* @return The resource content, or `null` if the resource is not found.
* @throws IllegalStateException If the server does not support resources.
*/
suspend fun readResource(request: ReadResourceRequest, options: RequestOptions? = null): ReadResourceResult? {
return request<ReadResourceResult>(request, options)
}
/**
* Subscribes to resource changes on the server.
*
* @param request The subscription request containing resource details.
* @param options Optional request options.
* @throws IllegalStateException If the server does not support resource subscriptions.
*/
suspend fun subscribeResource(request: SubscribeRequest, options: RequestOptions? = null): EmptyRequestResult {
return request<EmptyRequestResult>(request, options)
}
/**
* Unsubscribes from resource changes on the server.
*
* @param request The unsubscribe request containing resource details.
* @param options Optional request options.
* @throws IllegalStateException If the server does not support resource subscriptions.
*/
suspend fun unsubscribeResource(request: UnsubscribeRequest, options: RequestOptions? = null): EmptyRequestResult {
return request<EmptyRequestResult>(request, options)
}
/**
* Calls a tool on the server by name, passing the specified arguments.
*
* @param name The name of the tool to call.
* @param arguments A map of argument names to values for the tool.
* @param compatibility Whether to use compatibility mode for older protocol versions.
* @param options Optional request options.
* @return The result of the tool call, or `null` if none.
* @throws IllegalStateException If the server does not support tools.
*/
suspend fun callTool(
name: String,
arguments: Map<String, Any?>,
compatibility: Boolean = false,
options: RequestOptions? = null,
): CallToolResultBase? {
val jsonArguments = arguments.mapValues { (_, value) ->
when (value) {
is String -> JsonPrimitive(value)
is Number -> JsonPrimitive(value)
is Boolean -> JsonPrimitive(value)
is JsonElement -> value
null -> JsonNull
else -> JsonPrimitive(value.toString())
}
}
val request = CallToolRequest(
name = name,
arguments = JsonObject(jsonArguments)
)
return callTool(request, compatibility, options)
}
/**
* Calls a tool on the server using a [CallToolRequest] object.
*
* @param request The request object containing the tool name and arguments.
* @param compatibility Whether to use compatibility mode for older protocol versions.
* @param options Optional request options.
* @return The result of the tool call, or `null` if none.
* @throws IllegalStateException If the server does not support tools.
*/
suspend fun callTool(
request: CallToolRequest,
compatibility: Boolean = false,
options: RequestOptions? = null,
): CallToolResultBase? {
return if (compatibility) {
request<CompatibilityCallToolResult>(request, options)
} else {
request<CallToolResult>(request, options)
}
}
/**
* Lists all available tools on the server.
*
* @param request A request object for listing tools (usually empty).
* @param options Optional request options.
* @return The list of available tools, or `null` if none.
* @throws IllegalStateException If the server does not support tools.
*/
suspend fun listTools(
request: ListToolsRequest = ListToolsRequest(),
options: RequestOptions? = null,
): ListToolsResult? {
return request<ListToolsResult>(request, options)
}
/**
* Notifies the server that the list of roots has changed.
* Typically used if the client is managing some form of hierarchical structure.
*
* @throws IllegalStateException If the client or server does not support roots.
*/
suspend fun sendRootsListChanged() {
notification(RootsListChangedNotification())
}
}

138
src/main/kotlin/org/jetbrains/kotlinx/mcp/client/SSEClientTransport.kt Normal file
View file

@ -0,0 +1,138 @@
package org.jetbrains.kotlinx.mcp.client
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import io.ktor.client.*
import io.ktor.client.plugins.sse.*
import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.http.*
import kotlinx.coroutines.*
import kotlinx.serialization.encodeToString
import org.jetbrains.kotlinx.mcp.shared.McpJson
import org.jetbrains.kotlinx.mcp.shared.Transport
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.properties.Delegates
import kotlin.time.Duration
/**
* Client transport for SSE: this will connect to a server using Server-Sent Events for receiving
* messages and make separate POST requests for sending messages.
*/
class SSEClientTransport(
private val client: HttpClient,
private val urlString: String?,
private val reconnectionTime: Duration? = null,
private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
) : Transport {
private val scope by lazy {
CoroutineScope(session.coroutineContext + SupervisorJob())
}
private val initialized = AtomicBoolean(false)
private var session: ClientSSESession by Delegates.notNull()
private val endpoint = CompletableDeferred<String>()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
private var job: Job? = null
private val baseUrl by lazy {
session.call.request.url.toString().removeSuffix("/")
}
override suspend fun start() {
if (!initialized.compareAndSet(false, true)) {
error(
"SSEClientTransport already started! " +
"If using Client class, note that connect() calls start() automatically.",
)
}
session = urlString?.let {
client.sseSession(
urlString = it,
reconnectionTime = reconnectionTime,
block = requestBuilder,
)
} ?: client.sseSession(
reconnectionTime = reconnectionTime,
block = requestBuilder,
)
job = scope.launch(CoroutineName("SseMcpClientTransport.collect#${hashCode()}")) {
session.incoming.collect { event ->
when (event.event) {
"error" -> {
val e = IllegalStateException("SSE error: ${event.data}")
onError?.invoke(e)
throw e
}
"open" -> {
// The connection is open, but we need to wait for the endpoint to be received.
}
"endpoint" -> {
try {
val eventData = event.data ?: ""
// check url correctness
val maybeEndpoint = Url("$baseUrl/$eventData")
endpoint.complete(maybeEndpoint.toString())
} catch (e: Exception) {
onError?.invoke(e)
close()
error(e)
}
}
else -> {
try {
val message = McpJson.decodeFromString<JSONRPCMessage>(event.data ?: "")
onMessage?.invoke(message)
} catch (e: Exception) {
onError?.invoke(e)
}
}
}
}
}
endpoint.await()
}
@OptIn(ExperimentalCoroutinesApi::class)
override suspend fun send(message: JSONRPCMessage) {
if (!endpoint.isCompleted) {
error("Not connected")
}
try {
val response = client.post(endpoint.getCompleted()) {
headers.append(HttpHeaders.ContentType, ContentType.Application.Json)
setBody(McpJson.encodeToString(message))
}
if (!response.status.isSuccess()) {
val text = response.bodyAsText()
error("Error POSTing to endpoint (HTTP ${response.status}): $text")
}
} catch (e: Exception) {
onError?.invoke(e)
throw e
}
}
override suspend fun close() {
if (!initialized.get()) {
error("SSEClientTransport is not initialized!")
}
session.cancel()
onClose?.invoke()
job?.cancelAndJoin()
}
}

131
src/main/kotlin/org/jetbrains/kotlinx/mcp/client/StdioClientTransport.kt Normal file
View file

@ -0,0 +1,131 @@
package org.jetbrains.kotlinx.mcp.client
import io.github.oshai.kotlinlogging.KotlinLogging
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.consumeEach
import org.jetbrains.kotlinx.mcp.shared.ReadBuffer
import org.jetbrains.kotlinx.mcp.shared.Transport
import org.jetbrains.kotlinx.mcp.shared.serializeMessage
import java.io.InputStream
import java.io.OutputStream
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.coroutines.CoroutineContext
import kotlin.text.Charsets.UTF_8
/**
* A transport implementation for JSON-RPC communication that leverages standard input and output streams.
*
* This class reads from an input stream to process incoming JSON-RPC messages and writes JSON-RPC messages
* to an output stream.
*
* @param input The input stream where messages are received.
* @param output The output stream where messages are sent.
*/
class StdioClientTransport(
private val input: InputStream,
private val output: OutputStream
) : Transport {
private val logger = KotlinLogging.logger {}
private val ioCoroutineContext: CoroutineContext = Dispatchers.IO
private val scope by lazy {
CoroutineScope(ioCoroutineContext + SupervisorJob())
}
private var job: Job? = null
private var initialized = AtomicBoolean(false)
private val sendChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)
private val readBuffer = ReadBuffer()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
override suspend fun start() {
if (!initialized.compareAndSet(false, true)) {
error("StdioClientTransport already started!")
}
logger.debug { "Starting StdioClientTransport..." }
val outputStream = output.bufferedWriter(UTF_8)
job = scope.launch(CoroutineName("StdioClientTransport.IO#${hashCode()}")) {
val readJob = launch {
logger.debug { "Read coroutine started." }
try {
val buffer = ByteArray(8192)
while (isActive) {
val bytesRead = input.read(buffer)
if (bytesRead == -1) break
if (bytesRead > 0) {
readBuffer.append(buffer.copyOf(bytesRead))
processReadBuffer()
}
}
} catch (e: CancellationException) {
throw e
} catch (e: Exception) {
onError?.invoke(e)
logger.error(e) { "Error reading from input stream" }
} finally {
input.close()
}
}
val writeJob = launch {
logger.debug { "Write coroutine started." }
try {
sendChannel.consumeEach { message ->
val json = serializeMessage(message)
outputStream.write(json)
outputStream.flush()
}
} catch (e: Throwable) {
if (isActive) {
onError?.invoke(e)
logger.error(e) { "Error writing to output stream" }
}
} finally {
output.close()
}
}
readJob.join()
writeJob.cancelAndJoin()
onClose?.invoke()
}
}
override suspend fun send(message: JSONRPCMessage) {
if (!initialized.get()) {
error("Transport not started")
}
sendChannel.send(message)
}
override suspend fun close() {
if (!initialized.compareAndSet(true, false)) {
error("Transport is already closed")
}
job?.cancelAndJoin()
input.close()
output.close()
readBuffer.clear()
sendChannel.close()
onClose?.invoke()
}
private suspend fun processReadBuffer() {
while (true) {
val msg = readBuffer.readMessage() ?: break
try {
onMessage?.invoke(msg)
} catch (e: Throwable) {
onError?.invoke(e)
logger.error(e) { "Error processing message." }
}
}
}
}

35
src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketClientTransport.kt Normal file
View file

@ -0,0 +1,35 @@
package org.jetbrains.kotlinx.mcp.client
import io.ktor.client.*
import io.ktor.client.plugins.websocket.*
import io.ktor.client.request.*
import io.ktor.http.*
import io.ktor.websocket.*
import org.jetbrains.kotlinx.mcp.shared.MCP_SUBPROTOCOL
import org.jetbrains.kotlinx.mcp.shared.WebSocketMcpTransport
import kotlin.properties.Delegates
/**
* Client transport for WebSocket: this will connect to a server over the WebSocket protocol.
*/
class WebSocketClientTransport(
private val client: HttpClient,
private val urlString: String?,
private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
) : WebSocketMcpTransport() {
override var session: WebSocketSession by Delegates.notNull()
override suspend fun initializeSession() {
session = urlString?.let {
client.webSocketSession(it) {
requestBuilder()
header(HttpHeaders.SecWebSocketProtocol, MCP_SUBPROTOCOL)
}
} ?: client.webSocketSession {
requestBuilder()
header(HttpHeaders.SecWebSocketProtocol, MCP_SUBPROTOCOL)
}
}
}

41
src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketMcpKtorClientExtensions.kt Normal file
View file

@ -0,0 +1,41 @@
package org.jetbrains.kotlinx.mcp.client
import org.jetbrains.kotlinx.mcp.Implementation
import io.ktor.client.HttpClient
import io.ktor.client.request.HttpRequestBuilder
import org.jetbrains.kotlinx.mcp.shared.IMPLEMENTATION_NAME
import shared.LIB_VERSION
/**
* Returns a new WebSocket transport for the Model Context Protocol using the provided HttpClient.
*
* @param urlString Optional URL of the MCP server.
* @param requestBuilder Optional lambda to configure the HTTP request.
* @return A [WebSocketClientTransport] configured for MCP communication.
*/
fun HttpClient.mcpWebSocketTransport(
urlString: String? = null,
requestBuilder: HttpRequestBuilder.() -> Unit = {},
): WebSocketClientTransport = WebSocketClientTransport(this, urlString, requestBuilder)
/**
* Creates and connects an MCP client over WebSocket using the provided HttpClient.
*
* @param urlString Optional URL of the MCP server.
* @param requestBuilder Optional lambda to configure the HTTP request.
* @return A connected [Client] ready for MCP communication.
*/
suspend fun HttpClient.mcpWebSocket(
urlString: String? = null,
requestBuilder: HttpRequestBuilder.() -> Unit = {},
): Client {
val transport = mcpWebSocketTransport(urlString, requestBuilder)
val client = Client(
Implementation(
name = IMPLEMENTATION_NAME,
version = LIB_VERSION,
)
)
client.connect(transport)
return client
}

46
src/main/kotlin/org/jetbrains/kotlinx/mcp/client/sse.ktor.kt Normal file
View file

@ -0,0 +1,46 @@
package org.jetbrains.kotlinx.mcp.client
import org.jetbrains.kotlinx.mcp.Implementation
import io.ktor.client.HttpClient
import io.ktor.client.request.HttpRequestBuilder
import org.jetbrains.kotlinx.mcp.shared.IMPLEMENTATION_NAME
import shared.LIB_VERSION
import kotlin.time.Duration
/**
* Returns a new SSE transport for the Model Context Protocol using the provided HttpClient.
*
* @param urlString Optional URL of the MCP server.
* @param reconnectionTime Optional duration to wait before attempting to reconnect.
* @param requestBuilder Optional lambda to configure the HTTP request.
* @return A [SSEClientTransport] configured for MCP communication.
*/
fun HttpClient.mcpSseTransport(
urlString: String? = null,
reconnectionTime: Duration? = null,
requestBuilder: HttpRequestBuilder.() -> Unit = {},
): SSEClientTransport = SSEClientTransport(this, urlString, reconnectionTime, requestBuilder)
/**
* Creates and connects an MCP client over SSE using the provided HttpClient.
*
* @param urlString Optional URL of the MCP server.
* @param reconnectionTime Optional duration to wait before attempting to reconnect.
* @param requestBuilder Optional lambda to configure the HTTP request.
* @return A connected [Client] ready for MCP communication.
*/
suspend fun HttpClient.mcpSse(
urlString: String? = null,
reconnectionTime: Duration? = null,
requestBuilder: HttpRequestBuilder.() -> Unit = {},
): Client {
val transport = mcpSseTransport(urlString, reconnectionTime, requestBuilder)
val client = Client(
Implementation(
name = IMPLEMENTATION_NAME,
version = LIB_VERSION,
)
)
client.connect(transport)
return client
}

54
src/main/kotlin/org/jetbrains/kotlinx/mcp/server/McpKtorServerPlugin.kt Normal file
View file

@ -0,0 +1,54 @@
package org.jetbrains.kotlinx.mcp.server
import io.github.oshai.kotlinlogging.KotlinLogging
import io.ktor.http.*
import io.ktor.server.application.*
import io.ktor.server.response.*
import io.ktor.server.routing.*
import io.ktor.server.sse.*
import io.ktor.util.collections.*
private val logger = KotlinLogging.logger {}
/**
* Configures the Ktor Application to handle Model Context Protocol (MCP) over Server-Sent Events (SSE).
*/
fun Application.MCP(block: () -> Server) {
val servers = ConcurrentMap<String, Server>()
install(SSE)
routing {
sse("/sse") {
val transport = SSEServerTransport("/message", this)
logger.info { "New SSE connection established with sessionId: ${transport.sessionId}" }
val server = block()
servers[transport.sessionId] = server
logger.debug { "Server instance created and stored for sessionId: ${transport.sessionId}" }
server.onCloseCallback = {
logger.info { "Server connection closed for sessionId: ${transport.sessionId}" }
servers.remove(transport.sessionId)
}
server.connect(transport)
logger.debug { "Server connected to transport for sessionId: ${transport.sessionId}" }
}
post("/message") {
val sessionId: String = call.request.queryParameters["sessionId"]!!
logger.debug { "Received message for sessionId: $sessionId" }
val transport = servers[sessionId]?.transport as? SSEServerTransport
if (transport == null) {
logger.warn { "Session not found for sessionId: $sessionId" }
call.respond(HttpStatusCode.NotFound, "Session not found")
return@post
}
transport.handlePostMessage(call)
logger.trace { "Message handled for sessionId: $sessionId" }
}
}
}

124
src/main/kotlin/org/jetbrains/kotlinx/mcp/server/SSEServerTransport.kt Normal file
View file

@ -0,0 +1,124 @@
package org.jetbrains.kotlinx.mcp.server
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import io.ktor.http.*
import io.ktor.server.application.*
import io.ktor.server.request.*
import io.ktor.server.response.*
import io.ktor.server.sse.*
import kotlinx.coroutines.job
import kotlinx.serialization.encodeToString
import org.jetbrains.kotlinx.mcp.shared.McpJson
import org.jetbrains.kotlinx.mcp.shared.Transport
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.uuid.ExperimentalUuidApi
import kotlin.uuid.Uuid
internal const val SESSION_ID_PARAM = "sessionId"
/**
* 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`.
*/
class SSEServerTransport(
private val endpoint: String,
private val session: ServerSSESession,
) : Transport {
private val initialized = AtomicBoolean(false)
@OptIn(ExperimentalUuidApi::class)
val sessionId: String = Uuid.random().toString()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
/**
* 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(false, true)) {
throw 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.
*/
suspend fun handlePostMessage(call: ApplicationCall) {
if (!initialized.get()) {
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.
*/
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.get()) {
throw error("Not connected")
}
session.send(
event = "message",
data = McpJson.encodeToString(message),
)
}
}

571
src/main/kotlin/org/jetbrains/kotlinx/mcp/server/Server.kt Normal file
View file

@ -0,0 +1,571 @@
package org.jetbrains.kotlinx.mcp.server
import io.github.oshai.kotlinlogging.KotlinLogging
import kotlinx.coroutines.CompletableDeferred
import kotlinx.serialization.json.JsonObject
import org.jetbrains.kotlinx.mcp.*
import org.jetbrains.kotlinx.mcp.shared.Protocol
import org.jetbrains.kotlinx.mcp.shared.ProtocolOptions
import org.jetbrains.kotlinx.mcp.shared.RequestOptions
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.
*/
class ServerOptions(
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.
* @param onCloseCallback A callback invoked when the server connection closes.
*/
open class Server(
private val serverInfo: Implementation,
options: ServerOptions,
var onCloseCallback: (() -> Unit)? = null
) : Protocol<ServerRequest, ServerNotification, ServerResult>(options) {
/**
* The client's reported capabilities after initialization.
*/
var clientCapabilities: ClientCapabilities? = null
private set
/**
* The client's version information after initialization.
*/
var clientVersion: Implementation? = null
private set
private val capabilities: ServerCapabilities = options.capabilities
private val tools = mutableMapOf<String, RegisteredTool>()
private val prompts = mutableMapOf<String, RegisteredPrompt>()
private val resources = mutableMapOf<String, RegisteredResource>()
/**
* A callback invoked when the server has completed the initialization sequence.
* After initialization, the server is ready to handle requests.
*/
var onInitialized: (() -> Unit)? = null
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?.invoke()
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()
}
}
}
/**
* Called when the server connection is closing.
* Invokes [onCloseCallback] if set.
*/
override fun onclose() {
logger.info { "Server connection closing" }
onCloseCallback?.invoke()
}
/**
* Registers a single tool. This tool can then be called by the client.
*
* @param name The name of the tool.
* @param description A human-readable description of what the tool does.
* @param inputSchema The expected input schema for 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.
*/
fun addTool(
name: String,
description: String,
inputSchema: Tool.Input = Tool.Input(),
handler: suspend (CallToolRequest) -> CallToolResult
) {
if (capabilities.tools == null) {
logger.error { "Failed to add tool '$name': Server does not support tools capability" }
throw IllegalStateException("Server does not support tools capability. Enable it in ServerOptions.")
}
logger.info { "Registering tool: $name" }
tools[name] = RegisteredTool(Tool(name, description, inputSchema), 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.
*/
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" }
for (rt in toolsToAdd) {
logger.debug { "Registering tool: ${rt.tool.name}" }
tools[rt.tool.name] = rt
}
}
/**
* Registers a single prompt. The prompt can then be retrieved by the client.
*
* @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.
*/
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[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.
*/
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.
*/
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" }
for (rp in promptsToAdd) {
logger.debug { "Registering prompt: ${rp.prompt.name}" }
prompts[rp.prompt.name] = rp
}
}
/**
* Registers a single resource. The resource content can then be read by the client.
*
* @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.
*/
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[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.
*/
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" }
for (r in resourcesToAdd) {
logger.debug { "Registering resource: ${r.resource.name} (${r.resource.uri})" }
resources[r.resource.uri] = r
}
}
/**
* 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.
*/
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.
*/
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.
*/
suspend fun listRoots(
params: JsonObject = EmptyJsonObject,
options: RequestOptions? = null
): ListRootsResult {
logger.debug { "Listing roots with params: $params" }
return request<ListRootsResult>(ListRootsRequest(params), options)
}
/**
* Sends a logging message notification to the client.
*
* @param params The logging message notification parameters.
*/
suspend fun sendLoggingMessage(params: LoggingMessageNotification) {
logger.trace { "Sending logging message: ${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.
*/
suspend fun sendResourceUpdated(params: ResourceUpdatedNotification) {
logger.debug { "Sending resource updated notification for: ${params.uri}" }
notification(params)
}
/**
* Sends a notification to the client indicating that the list of resources has changed.
*/
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.
*/
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.
*/
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[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})")
}
}
"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.
*/
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.
*/
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.
*/
data class RegisteredResource(
val resource: Resource,
val readHandler: suspend (ReadResourceRequest) -> ReadResourceResult
)

118
src/main/kotlin/org/jetbrains/kotlinx/mcp/server/StdioServerTransport.kt Normal file
View file

@ -0,0 +1,118 @@
package org.jetbrains.kotlinx.mcp.server
import io.github.oshai.kotlinlogging.KotlinLogging
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import org.jetbrains.kotlinx.mcp.shared.ReadBuffer
import org.jetbrains.kotlinx.mcp.shared.Transport
import org.jetbrains.kotlinx.mcp.shared.serializeMessage
import java.io.BufferedInputStream
import java.io.PrintStream
import java.util.concurrent.atomic.AtomicBoolean
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.
*/
class StdioServerTransport(
private val inputStream: BufferedInputStream = BufferedInputStream(System.`in`),
private val outputStream: PrintStream = System.out
) : Transport {
private val logger = KotlinLogging.logger {}
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null
private val readBuffer = ReadBuffer()
private var initialized = AtomicBoolean(false)
private var readingJob: Job? = null
private val coroutineContext: CoroutineContext = Dispatchers.IO + SupervisorJob()
private val scope = CoroutineScope(coroutineContext)
private val readChannel = Channel<ByteArray>(Channel.UNLIMITED)
private val outputWriter = outputStream.bufferedWriter()
override suspend fun start() {
if (!initialized.compareAndSet(false, true)) {
error("StdioServerTransport already started!")
}
// Launch a coroutine to read from stdin
readingJob = scope.launch {
val buf = ByteArray(8192)
try {
while (isActive) {
val bytesRead = inputStream.read(buf)
if (bytesRead == -1) {
// EOF reached
break
}
if (bytesRead > 0) {
val chunk = buf.copyOf(bytesRead)
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)
}
}
}
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(true, false)) return
// Cancel reading job and close channel
readingJob?.cancel() // ToDO("was cancel and join")
readChannel.close()
readBuffer.clear()
onClose?.invoke()
}
override suspend fun send(message: JSONRPCMessage) {
val json = serializeMessage(message)
synchronized(outputWriter) {
// You may need to add Content-Length headers before the message if using the LSP framing protocol
outputWriter.write(json)
outputWriter.flush()
}
}
}

107
src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpKtorServerExtensions.kt Normal file
View file

@ -0,0 +1,107 @@
package org.jetbrains.kotlinx.mcp.server
import org.jetbrains.kotlinx.mcp.Implementation
import org.jetbrains.kotlinx.mcp.ServerCapabilities
import io.ktor.server.routing.*
import io.ktor.server.websocket.*
import org.jetbrains.kotlinx.mcp.shared.IMPLEMENTATION_NAME
import shared.LIB_VERSION
/**
* 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.
*/
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.
*/
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.
*/
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.
*/
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)
}

22
src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpServerTransport.kt Normal file
View file

@ -0,0 +1,22 @@
package org.jetbrains.kotlinx.mcp.server
import io.ktor.http.*
import io.ktor.server.websocket.*
import org.jetbrains.kotlinx.mcp.shared.MCP_SUBPROTOCOL
import org.jetbrains.kotlinx.mcp.shared.WebSocketMcpTransport
/**
* Server-side implementation of the MCP (Model Context Protocol) transport over WebSocket.
*
* @property session The WebSocket server session used for communication.
*/
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")
}
}
}

490
src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Protocol.kt Normal file
View file

@ -0,0 +1,490 @@
package org.jetbrains.kotlinx.mcp.shared
import org.jetbrains.kotlinx.mcp.CancelledNotification
import org.jetbrains.kotlinx.mcp.EmptyRequestResult
import org.jetbrains.kotlinx.mcp.ErrorCode
import org.jetbrains.kotlinx.mcp.JSONRPCError
import org.jetbrains.kotlinx.mcp.JSONRPCNotification
import org.jetbrains.kotlinx.mcp.JSONRPCRequest
import org.jetbrains.kotlinx.mcp.JSONRPCResponse
import org.jetbrains.kotlinx.mcp.McpError
import org.jetbrains.kotlinx.mcp.Method
import org.jetbrains.kotlinx.mcp.Notification
import org.jetbrains.kotlinx.mcp.PingRequest
import org.jetbrains.kotlinx.mcp.Progress
import org.jetbrains.kotlinx.mcp.ProgressNotification
import org.jetbrains.kotlinx.mcp.Request
import org.jetbrains.kotlinx.mcp.RequestResult
import io.github.oshai.kotlinlogging.KotlinLogging
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.withTimeout
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.*
import kotlinx.serialization.serializer
import org.jetbrains.kotlinx.mcp.fromJSON
import org.jetbrains.kotlinx.mcp.toJSON
import kotlin.collections.get
import kotlin.reflect.typeOf
import kotlin.time.Duration
import kotlin.time.Duration.Companion.milliseconds
private val LOGGER = KotlinLogging.logger { }
internal const val IMPLEMENTATION_NAME = "mcp-ktor"
/**
* Callback for progress notifications.
*/
typealias ProgressCallback = (Progress) -> Unit
@OptIn(ExperimentalSerializationApi::class)
@PublishedApi
internal val McpJson: Json by lazy {
Json {
ignoreUnknownKeys = true
encodeDefaults = true
isLenient = true
classDiscriminatorMode = ClassDiscriminatorMode.NONE
explicitNulls = false
}
}
/**
* Additional initialization options.
*/
open class ProtocolOptions(
/**
* Whether to restrict emitted requests to only those that the remote side has indicated
* that they can handle, through their advertised capabilities.
*
* Note that this DOES NOT affect checking of _local_ side capabilities, as it is
* considered a logic error to mis-specify those.
*
* Currently this defaults to false, for backwards compatibility with SDK versions
* that did not advertise capabilities correctly. In future, this will default to true.
*/
var enforceStrictCapabilities: Boolean = false,
var timeout: Duration = DEFAULT_REQUEST_TIMEOUT
)
/**
* The default request timeout.
*/
val DEFAULT_REQUEST_TIMEOUT: Duration = 60000.milliseconds
/**
* Options that can be given per request.
*/
data class RequestOptions(
/**
* If set, requests progress notifications from the remote end (if supported).
* When progress notifications are received, this callback will be invoked.
*/
val onProgress: ProgressCallback? = null,
/**
* A timeout for this request. If exceeded, an org.jetbrains.kotlinx.mcp.McpError with code `RequestTimeout`
* will be raised from request().
*
* If not specified, `DEFAULT_REQUEST_TIMEOUT` will be used as the timeout.
*/
val timeout: Duration = DEFAULT_REQUEST_TIMEOUT
)
/**
* Extra data given to request handlers.
*/
class RequestHandlerExtra()
internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) }
/**
* Implements MCP protocol framing on top of a pluggable transport, including
* features like request/response linking, notifications, and progress.
*/
abstract class Protocol<SendRequestT : Request, SendNotificationT : Notification, SendResultT : RequestResult>(
@PublishedApi internal val _options: ProtocolOptions?
) {
var transport: Transport? = null
private set
@PublishedApi
internal val requestHandlers: MutableMap<String, suspend (request: JSONRPCRequest, extra: RequestHandlerExtra) -> SendResultT?> =
mutableMapOf()
val notificationHandlers: MutableMap<String, suspend (notification: JSONRPCNotification) -> Unit> =
mutableMapOf()
@PublishedApi
internal val responseHandlers: MutableMap<Long, (response: JSONRPCResponse?, error: Exception?) -> Unit> =
mutableMapOf()
@PublishedApi
internal val progressHandlers: MutableMap<Long, ProgressCallback> = mutableMapOf()
/**
* Callback for when the connection is closed for any reason.
*
* This is invoked when close() is called as well.
*/
open fun onclose() {}
/**
* Callback for when an error occurs.
*
* Note that errors are not necessarily fatal they are used for reporting any kind of exceptional condition out of band.
*/
open fun onerror(error: Throwable) {}
/**
* A handler to invoke for any request types that do not have their own handler installed.
*/
var fallbackRequestHandler: (suspend (request: JSONRPCRequest, extra: RequestHandlerExtra) -> SendResultT?)? = null
/**
* A handler to invoke for any notification types that do not have their own handler installed.
*/
var fallbackNotificationHandler: (suspend (notification: JSONRPCNotification) -> Unit)? = null
init {
setNotificationHandler<ProgressNotification>(Method.Defined.NotificationsProgress) { notification ->
this.onProgress(notification)
COMPLETED
}
setRequestHandler<PingRequest>(Method.Defined.Ping) { request, _ ->
@Suppress("UNCHECKED_CAST")
EmptyRequestResult() as SendResultT
}
}
/**
* Attaches to the given transport, starts it, and starts listening for messages.
*
* The Protocol object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward.
*/
open suspend fun connect(transport: Transport) {
this.transport = transport
transport.onClose = {
this.onClose()
}
transport.onError = {
this.onError(it)
}
transport.onMessage = { message ->
when (message) {
is JSONRPCResponse -> onResponse(message, null)
is JSONRPCRequest -> onRequest(message)
is JSONRPCNotification -> onNotification(message)
is JSONRPCError -> onResponse(null, message)
}
}
return transport.start()
}
private fun onClose() {
val error = McpError(ErrorCode.Defined.ConnectionClosed.code, "Connection closed")
for (handler in responseHandlers.values) {
handler(null, error)
}
responseHandlers.clear()
progressHandlers.clear()
transport = null
onclose()
}
private fun onError(error: Throwable) {
onerror(error)
}
private suspend fun onNotification(notification: JSONRPCNotification) {
LOGGER.trace { "Received notification: ${notification.method}" }
val function = notificationHandlers[notification.method]
val property = fallbackNotificationHandler
val handler = function ?: property
if (handler == null) {
LOGGER.trace { "No handler found for notification: ${notification.method}" }
return
}
try {
handler(notification)
} catch (cause: Throwable) {
LOGGER.error(cause) { "Error handling notification: ${notification.method}" }
onError(cause)
}
}
private suspend fun onRequest(request: JSONRPCRequest) {
LOGGER.trace { "Received request: ${request.method} (id: ${request.id})" }
val handler = requestHandlers[request.method] ?: this.fallbackRequestHandler
if (handler === null) {
LOGGER.trace { "No handler found for request: ${request.method}" }
try {
transport!!.send(
JSONRPCResponse(
id = request.id,
error = JSONRPCError(
ErrorCode.Defined.MethodNotFound,
message = "Server does not support ${request.method}",
)
)
)
} catch (cause: Throwable) {
LOGGER.error(cause) { "Error sending method not found response" }
onError(cause)
}
return
}
try {
val result = handler(request, RequestHandlerExtra())
LOGGER.trace { "Request handled successfully: ${request.method} (id: ${request.id})" }
val response = JSONRPCResponse(
id = request.id,
result = result
)
transport!!.send(response)
} catch (cause: Throwable) {
LOGGER.error(cause) { "Error handling request: ${request.method} (id: ${request.id})" }
transport!!.send(JSONRPCResponse(
id = request.id,
error = JSONRPCError(
ErrorCode.Defined.InternalError,
message = cause.message ?: "Internal error",
)
))
}
}
private fun onProgress(notification: ProgressNotification) {
LOGGER.trace { "Received progress notification: token=${notification.progressToken}, progress=${notification.progress}/${notification.total}" }
val progress = notification.progress
val total = notification.total
val progressToken = notification.progressToken
val handler = this.progressHandlers[progressToken]
if (handler == null) {
val error = Error(
"Received a progress notification for an unknown token: ${McpJson.encodeToString(notification)}",
)
LOGGER.error { error.message }
this.onError(error)
return
}
handler.invoke(Progress(progress, total))
}
private fun onResponse(response: JSONRPCResponse?, error: JSONRPCError?) {
val messageId = response?.id
val handler = this.responseHandlers[messageId]
if (handler == null) {
this.onError(Error("Received a response for an unknown message ID: ${McpJson.encodeToString(response)}"))
return
}
this.responseHandlers.remove(messageId)
this.progressHandlers.remove(messageId)
if (response != null) {
handler(response, null)
} else {
check(error != null)
val error = McpError(
error.code.code,
error.message,
error.data,
)
handler(null, error)
}
}
/**
* Closes the connection.
*/
suspend fun close() {
transport?.close()
}
/**
* A method to check if a capability is supported by the remote side, for the given method to be called.
*
* This should be implemented by subclasses.
*/
abstract fun assertCapabilityForMethod(method: Method)
/**
* A method to check if a notification is supported by the local side, for the given method to be sent.
*
* This should be implemented by subclasses.
*/
protected abstract fun assertNotificationCapability(method: Method)
/**
* A method to check if a request handler is supported by the local side, for the given method to be handled.
*
* This should be implemented by subclasses.
*/
abstract fun assertRequestHandlerCapability(method: Method)
/**
* Sends a request and wait for a response.
*
* Do not use this method to emit notifications! Use notification() instead.
*/
suspend fun <T : RequestResult> request(
request: SendRequestT,
options: RequestOptions? = null,
): T {
LOGGER.trace { "Sending request: ${request.method}" }
val result = CompletableDeferred<T>()
val transport = this@Protocol.transport ?: throw Error("Not connected")
if (_options?.enforceStrictCapabilities == true) {
assertCapabilityForMethod(request.method)
}
val message = request.toJSON()
val messageId = message.id
if (options?.onProgress != null) {
LOGGER.trace { "Registering progress handler for request id: $messageId" }
progressHandlers[messageId] = options.onProgress
}
responseHandlers.set(messageId, { response, error ->
if (error != null) {
result.completeExceptionally(error)
return@set
}
if (response?.error != null) {
result.completeExceptionally(IllegalStateException(response.error.toString()))
return@set
}
try {
@Suppress("UNCHECKED_CAST")
result.complete(response!!.result as T)
} catch (error: Throwable) {
result.completeExceptionally(error)
}
})
val cancel: suspend (Throwable) -> Unit = { reason: Throwable ->
responseHandlers.remove(messageId)
progressHandlers.remove(messageId)
val notification = CancelledNotification(requestId = messageId, reason = reason.message ?: "Unknown")
val serialized = JSONRPCNotification(
notification.method.value,
params = McpJson.encodeToJsonElement(notification)
)
transport.send(serialized)
result.completeExceptionally(reason)
Unit
}
val timeout = options?.timeout ?: DEFAULT_REQUEST_TIMEOUT
try {
withTimeout(timeout) {
LOGGER.trace { "Sending request message with id: $messageId" }
this@Protocol.transport!!.send(message)
}
return result.await()
} catch (cause: TimeoutCancellationException) {
LOGGER.error { "Request timed out after ${timeout.inWholeMilliseconds}ms: ${request.method}" }
cancel(
McpError(
ErrorCode.Defined.RequestTimeout.code,
"Request timed out",
JsonObject(mutableMapOf("timeout" to JsonPrimitive(timeout.inWholeMilliseconds)))
),
)
result.cancel(cause)
throw cause
}
}
/**
* Emits a notification, which is a one-way message that does not expect a response.
*/
suspend fun notification(notification: SendNotificationT) {
LOGGER.trace { "Sending notification: ${notification.method}" }
val transport = this.transport ?: error("Not connected")
assertNotificationCapability(notification.method)
val message = JSONRPCNotification(
notification.method.value,
params = McpJson.encodeToJsonElement<Notification>(notification) as JsonObject,
)
transport.send(message)
}
/**
* Registers a handler to invoke when this protocol object receives a request with the given method.
*
* Note that this will replace any previous request handler for the same method.
*/
inline fun <reified T : Request> setRequestHandler(
method: Method,
noinline block: suspend (T, RequestHandlerExtra) -> SendResultT?
) {
assertRequestHandlerCapability(method)
val type = typeOf<T>()
val serializer = McpJson.serializersModule.serializer(type)
requestHandlers[method.value] = { request, extraHandler ->
val result = request.params?.let { McpJson.decodeFromJsonElement(serializer, it) }
val response = if (result != null) {
block(result as T, extraHandler)
} else {
EmptyRequestResult()
}
@Suppress("UNCHECKED_CAST")
response as SendResultT
}
}
/**
* Removes the request handler for the given method.
*/
fun removeRequestHandler(method: Method) {
requestHandlers.remove(method.value)
}
/**
* Registers a handler to invoke when this protocol object receives a notification with the given method.
*
* Note that this will replace any previous notification handler for the same method.
*/
fun <T : Notification> setNotificationHandler(method: Method, handler: (notification: T) -> Deferred<Unit>) {
this.notificationHandlers[method.value] = {
@Suppress("UNCHECKED_CAST")
handler(it.fromJSON() as T)
}
}
/**
* Removes the notification handler for the given method.
*/
fun removeNotificationHandler(method: Method) {
this.notificationHandlers.remove(method.value)
}
}

56
src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/ReadBuffer.kt Normal file
View file

@ -0,0 +1,56 @@
package org.jetbrains.kotlinx.mcp.shared
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import io.ktor.utils.io.core.*
import kotlinx.io.Buffer
import kotlinx.io.indexOf
import kotlinx.io.readString
import kotlinx.serialization.encodeToString
/**
* Buffers a continuous stdio stream into discrete JSON-RPC messages.
*/
class ReadBuffer {
private val buffer: Buffer = Buffer()
fun append(chunk: ByteArray) {
buffer.writeFully(chunk)
}
fun readMessage(): JSONRPCMessage? {
if (buffer.exhausted()) return null
var lfIndex = buffer.indexOf('\n'.code.toByte())
val line = when (lfIndex) {
-1L -> return null
0L -> {
buffer.skip(1)
""
}
else -> {
var skipBytes = 1
if (buffer[lfIndex - 1] == '\r'.code.toByte()) {
lfIndex -= 1
skipBytes += 1
}
val string = buffer.readString(lfIndex)
buffer.skip(skipBytes.toLong())
string
}
}
return deserializeMessage(line)
}
fun clear() {
buffer.clear()
}
}
internal fun deserializeMessage(line: String): JSONRPCMessage {
return McpJson.decodeFromString<JSONRPCMessage>(line)
}
internal fun serializeMessage(message: JSONRPCMessage): String {
return McpJson.encodeToString(message) + "\n"
}

48
src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Transport.kt Normal file
View file

@ -0,0 +1,48 @@
package org.jetbrains.kotlinx.mcp.shared
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
/**
* Describes the minimal contract for a MCP transport that a client or server can communicate over.
*/
interface Transport {
/**
* Starts processing messages on the transport, including any connection steps that might need to be taken.
*
* This method should only be called after callbacks are installed, or else messages may be lost.
*
* NOTE: This method should not be called explicitly when using Client, Server, or Protocol classes,
* as they will implicitly call start().
*/
suspend fun start()
/**
* Sends a JSON-RPC message (request or response).
*/
suspend fun send(message: JSONRPCMessage)
/**
* Closes the connection.
*/
suspend fun close()
/**
* Callback for when the connection is closed for any reason.
*
* This should be invoked when close() is called as well.
*/
var onClose: (() -> Unit)?
/**
* Callback for when an error occurs.
*
* Note that errors are not necessarily fatal; they are used for reporting any kind of
* exceptional condition out of band.
*/
var onError: ((Throwable) -> Unit)?
/**
* Callback for when a message (request or response) is received over the connection.
*/
var onMessage: (suspend ((JSONRPCMessage) -> Unit))?
}

104
src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/WebSocketMcpTransport.kt Normal file
View file

@ -0,0 +1,104 @@
package org.jetbrains.kotlinx.mcp.shared
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import io.ktor.websocket.Frame
import io.ktor.websocket.WebSocketSession
import io.ktor.websocket.close
import io.ktor.websocket.readText
import kotlinx.coroutines.CoroutineName
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.InternalCoroutinesApi
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.job
import kotlinx.coroutines.launch
import kotlinx.serialization.encodeToString
import java.util.concurrent.atomic.AtomicBoolean
internal const val MCP_SUBPROTOCOL = "mcp"
/**
* Abstract class representing a WebSocket transport for the Model Context Protocol (MCP).
* Handles communication over a WebSocket session.
*/
abstract class WebSocketMcpTransport : Transport {
private val scope by lazy {
CoroutineScope(session.coroutineContext + SupervisorJob())
}
private val initialized = AtomicBoolean(false)
/**
* The WebSocket session used for communication.
*/
protected abstract val session: WebSocketSession
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
/**
* Initializes the WebSocket session
*/
protected abstract suspend fun initializeSession()
override suspend fun start() {
if (!initialized.compareAndSet(false, true)) {
error(
"WebSocketClientTransport already started! " +
"If using Client class, note that connect() calls start() automatically.",
)
}
initializeSession()
scope.launch(CoroutineName("WebSocketMcpTransport.collect#${hashCode()}")) {
while (true) {
val message = try {
session.incoming.receive()
} catch (_: ClosedReceiveChannelException) {
return@launch
}
if (message !is Frame.Text) {
val e = IllegalArgumentException("Expected text frame, got ${message::class.simpleName}: $message")
onError?.invoke(e)
throw e
}
try {
val message = McpJson.decodeFromString<JSONRPCMessage>(message.readText())
onMessage?.invoke(message)
} catch (e: Exception) {
onError?.invoke(e)
throw e
}
}
}
@OptIn(InternalCoroutinesApi::class)
session.coroutineContext.job.invokeOnCompletion {
if (it != null) {
onError?.invoke(it)
} else {
onClose?.invoke()
}
}
}
override suspend fun send(message: JSONRPCMessage) {
if (!initialized.get()) {
error("Not connected")
}
session.outgoing.send(Frame.Text(McpJson.encodeToString(message)))
}
override suspend fun close() {
if (!initialized.get()) {
error("Not connected")
}
session.close()
session.coroutineContext.job.join()
}
}

1448
src/main/kotlin/org/jetbrains/kotlinx/mcp/types.kt Normal file
View file

File diff suppressed because it is too large Load diff

386
src/main/kotlin/org/jetbrains/kotlinx/mcp/types.util.kt Normal file
View file

@ -0,0 +1,386 @@
package org.jetbrains.kotlinx.mcp
import org.jetbrains.kotlinx.mcp.LoggingMessageNotification.SetLevelRequest
import kotlinx.serialization.DeserializationStrategy
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.descriptors.PrimitiveKind
import kotlinx.serialization.descriptors.PrimitiveSerialDescriptor
import kotlinx.serialization.descriptors.SerialDescriptor
import kotlinx.serialization.encoding.Decoder
import kotlinx.serialization.encoding.Encoder
import kotlinx.serialization.json.*
internal object ErrorCodeSerializer : KSerializer<ErrorCode> {
override val descriptor: SerialDescriptor =
PrimitiveSerialDescriptor("org.jetbrains.kotlinx.mcp.ErrorCode", PrimitiveKind.STRING)
@OptIn(ExperimentalSerializationApi::class)
override fun serialize(encoder: Encoder, value: ErrorCode) {
encoder.encodeInt(value.code)
}
override fun deserialize(decoder: Decoder): ErrorCode {
val decodedString = decoder.decodeInt()
return ErrorCode.Defined.entries.firstOrNull { it.code == decodedString }
?: ErrorCode.Unknown(decodedString)
}
}
internal object RequestMethodSerializer : KSerializer<Method> {
override val descriptor: SerialDescriptor =
PrimitiveSerialDescriptor("org.jetbrains.kotlinx.mcp.Method", PrimitiveKind.STRING)
@OptIn(ExperimentalSerializationApi::class)
override fun serialize(encoder: Encoder, value: Method) {
encoder.encodeString(value.value)
}
override fun deserialize(decoder: Decoder): Method {
val decodedString = decoder.decodeString()
return Method.Defined.entries.firstOrNull { it.value == decodedString }
?: Method.Custom(decodedString)
}
}
internal object StopReasonSerializer : KSerializer<StopReason> {
override val descriptor: SerialDescriptor =
PrimitiveSerialDescriptor("org.jetbrains.kotlinx.mcp.StopReason", PrimitiveKind.STRING)
@OptIn(ExperimentalSerializationApi::class)
override fun serialize(encoder: Encoder, value: StopReason) {
encoder.encodeString(value.value)
}
override fun deserialize(decoder: Decoder): StopReason {
val decodedString = decoder.decodeString()
return when (decodedString) {
StopReason.StopSequence.value -> StopReason.StopSequence
StopReason.MaxTokens.value -> StopReason.MaxTokens
StopReason.EndTurn.value -> StopReason.EndTurn
else -> StopReason.Other(decodedString)
}
}
}
internal object ReferencePolymorphicSerializer : JsonContentPolymorphicSerializer<Reference>(Reference::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<Reference> {
return when (element.jsonObject.getValue("type").jsonPrimitive.content) {
ResourceReference.TYPE -> ResourceReference.serializer()
PromptReference.TYPE -> PromptReference.serializer()
else -> UnknownReference.serializer()
}
}
}
internal object PromptMessageContentPolymorphicSerializer :
JsonContentPolymorphicSerializer<PromptMessageContent>(PromptMessageContent::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<PromptMessageContent> {
return when (element.jsonObject.getValue("type").jsonPrimitive.content) {
ImageContent.TYPE -> ImageContent.serializer()
TextContent.TYPE -> TextContent.serializer()
EmbeddedResource.TYPE -> EmbeddedResource.serializer()
else -> UnknownContent.serializer()
}
}
}
internal object PromptMessageContentTextOrImagePolymorphicSerializer :
JsonContentPolymorphicSerializer<PromptMessageContentTextOrImage>(PromptMessageContentTextOrImage::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<PromptMessageContentTextOrImage> {
return when (element.jsonObject.getValue("type").jsonPrimitive.content) {
ImageContent.TYPE -> ImageContent.serializer()
TextContent.TYPE -> TextContent.serializer()
else -> UnknownContent.serializer()
}
}
}
internal object ResourceContentsPolymorphicSerializer :
JsonContentPolymorphicSerializer<ResourceContents>(ResourceContents::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<ResourceContents> {
val jsonObject = element.jsonObject
return when {
jsonObject.contains("text") -> TextResourceContents.serializer()
jsonObject.contains("blob") -> BlobResourceContents.serializer()
else -> UnknownResourceContents.serializer()
}
}
}
internal fun selectRequestDeserializer(method: String): DeserializationStrategy<Request> {
selectClientRequestDeserializer(method)?.let { return it }
selectServerRequestDeserializer(method)?.let { return it }
return CustomRequest.serializer()
}
internal fun selectClientRequestDeserializer(method: String): DeserializationStrategy<ClientRequest>? {
return when (method) {
Method.Defined.Ping.value -> PingRequest.serializer()
Method.Defined.Initialize.value -> InitializeRequest.serializer()
Method.Defined.CompletionComplete.value -> CompleteRequest.serializer()
Method.Defined.LoggingSetLevel.value -> SetLevelRequest.serializer()
Method.Defined.PromptsGet.value -> GetPromptRequest.serializer()
Method.Defined.PromptsList.value -> ListPromptsRequest.serializer()
Method.Defined.ResourcesList.value -> ListResourcesRequest.serializer()
Method.Defined.ResourcesTemplatesList.value -> ListResourceTemplatesRequest.serializer()
Method.Defined.ResourcesRead.value -> ReadResourceRequest.serializer()
Method.Defined.ResourcesSubscribe.value -> SubscribeRequest.serializer()
Method.Defined.ResourcesUnsubscribe.value -> UnsubscribeRequest.serializer()
Method.Defined.ToolsCall.value -> CallToolRequest.serializer()
Method.Defined.ToolsList.value -> ListToolsRequest.serializer()
else -> null
}
}
//internal object ClientRequestPolymorphicSerializer :
// JsonContentPolymorphicSerializer<org.jetbrains.kotlinx.mcp.ClientRequest>(org.jetbrains.kotlinx.mcp.ClientRequest::class) {
// override fun selectDeserializer(element: JsonElement): DeserializationStrategy<org.jetbrains.kotlinx.mcp.ClientRequest> {
// val method = element.jsonObject.getOrDefault("method", null)?.jsonPrimitive?.content
// ?: error("No method in $element")
//
// return org.jetbrains.kotlinx.mcp.selectClientRequestDeserializer(method)
// ?: org.jetbrains.kotlinx.mcp.UnknownMethodRequestOrNotification.serializer()
// }
//}
private fun selectClientNotificationDeserializer(element: JsonElement): DeserializationStrategy<ClientNotification>? {
return when (element.jsonObject.getValue("method").jsonPrimitive.content) {
Method.Defined.NotificationsCancelled.value -> CancelledNotification.serializer()
Method.Defined.NotificationsProgress.value -> ProgressNotification.serializer()
Method.Defined.NotificationsInitialized.value -> InitializedNotification.serializer()
Method.Defined.NotificationsRootsListChanged.value -> RootsListChangedNotification.serializer()
else -> null
}
}
internal object ClientNotificationPolymorphicSerializer :
JsonContentPolymorphicSerializer<ClientNotification>(ClientNotification::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<ClientNotification> {
return selectClientNotificationDeserializer(element)
?: UnknownMethodRequestOrNotification.serializer()
}
}
internal fun selectServerRequestDeserializer(method: String): DeserializationStrategy<ServerRequest>? {
return when (method) {
Method.Defined.Ping.value -> PingRequest.serializer()
Method.Defined.SamplingCreateMessage.value -> CreateMessageRequest.serializer()
Method.Defined.RootsList.value -> ListRootsRequest.serializer()
else -> null
}
}
//internal object ServerRequestPolymorphicSerializer :
// JsonContentPolymorphicSerializer<org.jetbrains.kotlinx.mcp.ServerRequest>(org.jetbrains.kotlinx.mcp.ServerRequest::class) {
// override fun selectDeserializer(element: JsonElement): DeserializationStrategy<org.jetbrains.kotlinx.mcp.ServerRequest> {
// return org.jetbrains.kotlinx.mcp.selectServerRequestDeserializer(element)
// ?: org.jetbrains.kotlinx.mcp.UnknownMethodRequestOrNotification.serializer()
// }
//}
internal fun selectServerNotificationDeserializer(element: JsonElement): DeserializationStrategy<ServerNotification>? {
return when (element.jsonObject.getValue("method").jsonPrimitive.content) {
Method.Defined.NotificationsCancelled.value -> CancelledNotification.serializer()
Method.Defined.NotificationsProgress.value -> ProgressNotification.serializer()
Method.Defined.NotificationsMessage.value -> LoggingMessageNotification.serializer()
Method.Defined.NotificationsResourcesUpdated.value -> ResourceUpdatedNotification.serializer()
Method.Defined.NotificationsResourcesListChanged.value -> ResourceListChangedNotification.serializer()
Method.Defined.ToolsList.value -> ToolListChangedNotification.serializer()
Method.Defined.PromptsList.value -> PromptListChangedNotification.serializer()
else -> null
}
}
internal object ServerNotificationPolymorphicSerializer :
JsonContentPolymorphicSerializer<ServerNotification>(ServerNotification::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<ServerNotification> {
return selectServerNotificationDeserializer(element)
?: UnknownMethodRequestOrNotification.serializer()
}
}
internal object NotificationPolymorphicSerializer :
JsonContentPolymorphicSerializer<Notification>(Notification::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<Notification> {
return selectClientNotificationDeserializer(element)
?: selectServerNotificationDeserializer(element)
?: UnknownMethodRequestOrNotification.serializer()
}
}
internal object RequestPolymorphicSerializer :
JsonContentPolymorphicSerializer<Request>(Request::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<Request> {
val method = element.jsonObject.getOrDefault("method", null)?.jsonPrimitive?.content ?: run {
System.err.println("No method in $element")
Throwable().printStackTrace(System.err)
error("No method in $element")
}
return selectClientRequestDeserializer(method)
?: selectServerRequestDeserializer(method)
?: UnknownMethodRequestOrNotification.serializer()
}
}
/**
* Server messages schemas
*
* We deserialize by unique keys
*
* ```
* org.jetbrains.kotlinx.mcp.CallToolResult {
* content: org.jetbrains.kotlinx.mcp.PromptMessageContent {
* type: String,
* }
* isError: Boolean?,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.CompatibilityCallToolResult {
* content: org.jetbrains.kotlinx.mcp.PromptMessageContent {
* type: String,
* },
* isError: Boolean?,
* _meta: JsonObject,
* toolResult: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.CompleteResult {
* completion: Completion {
* values: Array<String>,
* total: Int?,
* hasMore: Boolean?,
* }
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.GetPromptResult {
* description: String?,
* messages: Array<org.jetbrains.kotlinx.mcp.PromptMessage>,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.InitializeResult {
* protocolVersion: String,
* capabilities: org.jetbrains.kotlinx.mcp.ServerCapabilities,
* serverInfo: org.jetbrains.kotlinx.mcp.Implementation,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.ListPromptsResult {
* prompts: Array<org.jetbrains.kotlinx.mcp.Prompt>,
* nextCursor: org.jetbrains.kotlinx.mcp.Cursor?,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.ListResourceTemplatesResult {
* resourceTemplates: Array<org.jetbrains.kotlinx.mcp.ResourceTemplate>,
* nextCursor: org.jetbrains.kotlinx.mcp.Cursor?,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.ListResourcesResult {
* resources: Array<org.jetbrains.kotlinx.mcp.Resource>,
* nextCursor: org.jetbrains.kotlinx.mcp.Cursor?,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.ListToolsResult {
* tools: Array<org.jetbrains.kotlinx.mcp.Tool>,
* nextCursor: org.jetbrains.kotlinx.mcp.Cursor?,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.ReadResourceResult {
* contents: Array<org.jetbrains.kotlinx.mcp.ResourceContents>,
* _meta: JsonObject,
* }
* ```
*/
private fun selectServerResultDeserializer(element: JsonElement): DeserializationStrategy<ServerResult>? {
val jsonObject = element.jsonObject
return when {
jsonObject.contains("tools") -> ListToolsResult.serializer()
jsonObject.contains("resources") -> ListResourcesResult.serializer()
jsonObject.contains("resourceTemplates") -> ListResourceTemplatesResult.serializer()
jsonObject.contains("prompts") -> ListPromptsResult.serializer()
jsonObject.contains("capabilities") -> InitializeResult.serializer()
jsonObject.contains("description") -> GetPromptResult.serializer()
jsonObject.contains("completion") -> CompleteResult.serializer()
jsonObject.contains("toolResult") -> CompatibilityCallToolResult.serializer()
jsonObject.contains("contents") -> ReadResourceResult.serializer()
jsonObject.contains("content") -> CallToolResult.serializer()
else -> null
}
}
/**
* Client messages schemas
*
* We deserialize by unique keys
*
* ```
* org.jetbrains.kotlinx.mcp.CreateMessageResult {
* model: String,
* stopReason: org.jetbrains.kotlinx.mcp.StopReason?,
* role: org.jetbrains.kotlinx.mcp.Role,
* content: org.jetbrains.kotlinx.mcp.PromptMessageContentTextOrImage,
* _meta: JsonObject,
* }
*
* org.jetbrains.kotlinx.mcp.ListRootsResult {
* roots: Array<org.jetbrains.kotlinx.mcp.Root>,
* _meta: JsonObject,
* }
*```
*/
private fun selectClientResultDeserializer(element: JsonElement): DeserializationStrategy<ClientResult>? {
val jsonObject = element.jsonObject
return when {
jsonObject.contains("model") -> CreateMessageResult.serializer()
jsonObject.contains("roots") -> ListRootsResult.serializer()
else -> null
}
}
internal object ServerResultPolymorphicSerializer :
JsonContentPolymorphicSerializer<ServerResult>(ServerResult::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<ServerResult> {
return selectServerResultDeserializer(element)
?: EmptyRequestResult.serializer()
}
}
internal object ClientResultPolymorphicSerializer :
JsonContentPolymorphicSerializer<ClientResult>(ClientResult::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<ClientResult> {
return selectClientResultDeserializer(element)
?: EmptyRequestResult.serializer()
}
}
internal object RequestResultPolymorphicSerializer :
JsonContentPolymorphicSerializer<RequestResult>(RequestResult::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<RequestResult> {
return selectClientResultDeserializer(element)
?: selectServerResultDeserializer(element)
?: EmptyRequestResult.serializer()
}
}
internal object JSONRPCMessagePolymorphicSerializer :
JsonContentPolymorphicSerializer<JSONRPCMessage>(JSONRPCMessage::class) {
override fun selectDeserializer(element: JsonElement): DeserializationStrategy<JSONRPCMessage> {
val jsonObject = element.jsonObject
return when {
jsonObject.contains("message") -> JSONRPCError.serializer()
!jsonObject.contains("method") -> JSONRPCResponse.serializer()
jsonObject.contains("id") -> JSONRPCRequest.serializer()
else -> JSONRPCNotification.serializer()
}
}
}
internal val EmptyJsonObject = JsonObject(emptyMap())

58
src/test/kotlin/InMemoryTransport.kt Normal file
View file

@ -0,0 +1,58 @@
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import org.jetbrains.kotlinx.mcp.shared.Transport
/**
* In-memory transport for creating clients and servers that talk to each other within the same process.
*/
class InMemoryTransport : Transport {
private val scope = CoroutineScope(SupervisorJob())
private var otherTransport: InMemoryTransport? = null
private val messageQueue: MutableList<JSONRPCMessage> = mutableListOf()
override var onClose: (() -> Unit)? = null
override var onError: ((Throwable) -> Unit)? = null
override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
/**
* Creates a pair of linked in-memory transports that can communicate with each other.
* One should be passed to a Client and one to a Server.
*/
companion object {
fun createLinkedPair(): Pair<InMemoryTransport, InMemoryTransport> {
val clientTransport = InMemoryTransport()
val serverTransport = InMemoryTransport()
clientTransport.otherTransport = serverTransport
serverTransport.otherTransport = clientTransport
return Pair(clientTransport, serverTransport)
}
}
override suspend fun start() {
// Process any messages that were queued before start was called
while (messageQueue.isNotEmpty()) {
messageQueue.removeFirstOrNull()?.let { message ->
onMessage?.invoke(message) // todo?
}
}
}
override suspend fun close() {
val other = otherTransport
otherTransport = null
other?.close()
onClose?.invoke()
}
override suspend fun send(message: JSONRPCMessage) {
val other = otherTransport ?: throw IllegalStateException("Not connected")
if (other.onMessage != null) {
other.onMessage?.invoke(message) // todo?
} else {
other.messageQueue.add(message)
}
}
}

63
src/test/kotlin/client/BaseTransportTest.kt Normal file
View file

@ -0,0 +1,63 @@
package client
import org.jetbrains.kotlinx.mcp.InitializedNotification
import org.jetbrains.kotlinx.mcp.JSONRPCMessage
import org.jetbrains.kotlinx.mcp.PingRequest
import kotlinx.coroutines.CompletableDeferred
import org.jetbrains.kotlinx.mcp.shared.Transport
import org.jetbrains.kotlinx.mcp.toJSON
import kotlin.test.assertEquals
import kotlin.test.assertFalse
import kotlin.test.assertTrue
import kotlin.test.fail
abstract class BaseTransportTest {
protected suspend fun testClientOpenClose(client: Transport) {
client.onError = { error ->
fail("Unexpected error: $error")
}
var didClose = false
client.onClose = { didClose = true }
client.start()
assertFalse(didClose, "Transport should not be closed immediately after start")
client.close()
assertTrue(didClose, "Transport should be closed after close() call")
}
protected suspend fun testClientRead(client: Transport) {
client.onError = { error ->
error.printStackTrace()
fail("Unexpected error: $error")
}
val messages = listOf<JSONRPCMessage>(
PingRequest().toJSON(),
InitializedNotification().toJSON()
)
val readMessages = mutableListOf<JSONRPCMessage>()
val finished = CompletableDeferred<Unit>()
client.onMessage = { message ->
readMessages.add(message)
if (message == messages.last()) {
finished.complete(Unit)
}
}
client.start()
for (message in messages) {
client.send(message)
}
finished.await()
assertEquals(messages, readMessages, "Assert messages received")
client.close()
}
}

Some files were not shown because too many files have changed in this diff Show more