diff --git a/.gitignore b/.gitignore index b1dff0d..997813f 100644 --- a/.gitignore +++ b/.gitignore @@ -42,4 +42,8 @@ bin/ .vscode/ ### Mac OS ### -.DS_Store \ No newline at end of file +.DS_Store + +### Node.js ### +node_modules +dist diff --git a/.idea/gradle.xml b/.idea/gradle.xml index 35cedcc..2e197de 100644 --- a/.idea/gradle.xml +++ b/.idea/gradle.xml @@ -1,11 +1,17 @@ + diff --git a/.idea/jsLibraryMappings.xml b/.idea/jsLibraryMappings.xml new file mode 100644 index 0000000..c93c655 --- /dev/null +++ b/.idea/jsLibraryMappings.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/.idea/kotlinc.xml b/.idea/kotlinc.xml new file mode 100644 index 0000000..c224ad5 --- /dev/null +++ b/.idea/kotlinc.xml @@ -0,0 +1,6 @@ + + + + + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml index 3950f45..6f735e5 100644 --- a/.idea/misc.xml +++ b/.idea/misc.xml @@ -1,10 +1,17 @@ + + + + + + + - \ No newline at end of file + diff --git a/.idea/uiDesigner.xml b/.idea/uiDesigner.xml new file mode 100644 index 0000000..2b63946 --- /dev/null +++ b/.idea/uiDesigner.xml @@ -0,0 +1,124 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..59a00a2 --- /dev/null +++ b/README.md @@ -0,0 +1,141 @@ +# MCP Kotlin SDK + +Kotlin implementation of the [Model Context Protocol](https://modelcontextprotocol.io) (MCP), providing both client and server capabilities for integrating with LLM surfaces. + +## Overview + +The Model Context Protocol allows applications to provide context for LLMs in a standardized way, separating the concerns of providing context from the actual LLM interaction. This Kotlin SDK implements the full MCP specification, making it easy to: + +- Build MCP clients that can connect to any MCP server +- Create MCP servers that expose resources, prompts and tools +- Use standard transports like stdio, SSE, and WebSocket +- Handle all MCP protocol messages and lifecycle events + +## Installation + +Add the JitPack repository to your build file: + +```kotlin +repositories { + maven { url = uri("https://jitpack.io") } +} +``` + +Add the dependency: + +```kotlin +dependencies { + implementation("com.github.e5l:mcp-kotlin-sdk:main-SNAPSHOT") +} +``` + +## Quick Start + +### Creating a Client + +```kotlin +import org.jetbrains.kotlinx.mcp.client.Client +import org.jetbrains.kotlinx.mcp.client.StdioClientTransport +import org.jetbrains.kotlinx.mcp.Implementation + +val client = Client( + clientInfo = Implementation( + name = "example-client", + version = "1.0.0" + ) +) + +val transport = StdioClientTransport( + inputStream = processInputStream, + outputStream = processOutputStream +) + +// Connect to server +client.connect(transport) + +// List available resources +val resources = client.listResources() + +// Read a specific resource +val resourceContent = client.readResource( + ReadResourceRequest(uri = "file:///example.txt") +) +``` + +### Creating a Server + +```kotlin +import org.jetbrains.kotlinx.mcp.server.Server +import org.jetbrains.kotlinx.mcp.server.ServerOptions +import org.jetbrains.kotlinx.mcp.server.StdioServerTransport +import org.jetbrains.kotlinx.mcp.ServerCapabilities + +val server = Server( + serverInfo = Implementation( + name = "example-server", + version = "1.0.0" + ), + options = ServerOptions( + capabilities = ServerCapabilities( + resources = ServerCapabilities.Resources( + subscribe = true, + listChanged = true + ) + ) + ) +) + +// Add a resource +server.addResource( + uri = "file:///example.txt", + name = "Example Resource", + description = "An example text file", + mimeType = "text/plain" +) { request -> + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "This is the content of the example resource.", + uri = request.uri, + mimeType = "text/plain" + ) + ) + ) +} + +// Start server with stdio transport +val transport = StdioServerTransport() +server.connect(transport) +``` + +### Using SSE Transport + +```kotlin +import io.ktor.server.application.* +import org.jetbrains.kotlinx.mcp.server.MCP + +fun Application.module() { + MCP { + Server( + serverInfo = Implementation( + name = "example-sse-server", + version = "1.0.0" + ), + options = ServerOptions( + capabilities = ServerCapabilities( + prompts = ServerCapabilities.Prompts(listChanged = null), + resources = ServerCapabilities.Resources(subscribe = null, listChanged = null) + ) + ) + ) + } +} +``` + +## Contributing + +Issues and pull requests are welcome on GitHub at https://github.com/e5l/mcp-kotlin-sdk. + +## License + +This project is licensed under the MIT License—see the [LICENSE](LICENSE) file for details. diff --git a/build.gradle.kts b/build.gradle.kts index f05c9e2..fddbfb3 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -1,5 +1,9 @@ +import org.jetbrains.dokka.gradle.engine.parameters.VisibilityModifier + plugins { - kotlin("jvm") version "2.0.21" + alias(libs.plugins.kotlin.jvm) + alias(libs.plugins.kotlin.serialization) + alias(libs.plugins.dokka) } group = "org.jetbrains.mcp" @@ -10,12 +14,72 @@ repositories { } dependencies { - testImplementation(kotlin("test")) + implementation(libs.kotlinx.serialization.json) + implementation(libs.ktor.client.apache) + implementation(libs.ktor.client.cio) + implementation(libs.ktor.server.cio) + implementation(libs.ktor.server.sse) + implementation(libs.ktor.server.websockets) + implementation(libs.kotlin.logging) + + testImplementation(libs.kotlin.test) + testImplementation(libs.mockk) + testImplementation(libs.ktor.server.test.host) + testImplementation(libs.kotlinx.coroutines.test) + testImplementation(libs.kotlinx.coroutines.debug) + } tasks.test { useJUnitPlatform() } + +abstract class GenerateLibVersionTask @Inject constructor( + @get:Input val libVersion: String, + @get:OutputDirectory val sourcesDir: File +) : DefaultTask() { + @TaskAction + fun generate() { + val sourceFile = File(sourcesDir, "LibVersion.kt") + + sourceFile.writeText( + """ + package shared + + const val LIB_VERSION = "$libVersion" + + """.trimIndent() + ) + } +} + +dokka { + moduleName.set("MCP Kotlin SDK") + + dokkaSourceSets.main { + sourceLink { + localDirectory.set(file("src/main/kotlin")) + remoteUrl("https://github.com/e5l/mcp-kotlin-sdk") + remoteLineSuffix.set("#L") + documentedVisibilities(VisibilityModifier.Public) + } + } + dokkaPublications.html { + outputDirectory.set(project.layout.projectDirectory.dir("docs")) + } +} + +val sourcesDir = File(project.layout.buildDirectory.asFile.get(), "generated-sources/libVersion") + +val generateLibVersionTask = + tasks.register("generateLibVersion", version.toString(), sourcesDir) + kotlin { jvmToolchain(21) -} \ No newline at end of file + + sourceSets { + main { + kotlin.srcDir(generateLibVersionTask.map { it.sourcesDir }) + } + } +} diff --git a/gradle.properties b/gradle.properties index 7fc6f1f..94704ad 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1 +1,4 @@ kotlin.code.style=official + +org.jetbrains.dokka.experimental.gradle.pluginMode=V2Enabled +org.jetbrains.dokka.experimental.gradle.pluginMode.noWarn=true diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml new file mode 100644 index 0000000..33d2411 --- /dev/null +++ b/gradle/libs.versions.toml @@ -0,0 +1,37 @@ +[versions] +# plugins version +kotlin = "2.0.21" +dokka = "2.0.0-Beta" + +# libraries version +serialization = "1.7.3" +coroutines = "1.9.0" +ktor = "3.0.2" +mockk = "1.13.13" +logging = "7.0.0" + +[libraries] +# Kotlinx libraries +kotlinx-serialization-json = { group = "org.jetbrains.kotlinx", name = "kotlinx-serialization-json", version.ref = "serialization" } +kotlin-logging = { group = "io.github.oshai", name = "kotlin-logging-jvm", version.ref = "logging" } + +# Ktor +ktor-client-apache = { group = "io.ktor", name = "ktor-client-apache", version.ref = "ktor" } +ktor-client-cio = { group = "io.ktor", name = "ktor-client-cio", version.ref = "ktor" } +ktor-server-sse = { group = "io.ktor", name = "ktor-server-sse", version.ref = "ktor" } +ktor-server-websockets = { group = "io.ktor", name = "ktor-server-websockets", version.ref = "ktor" } +ktor-server-cio = { group = "io.ktor", name = "ktor-server-cio", version.ref = "ktor" } + +# Testing +kotlin-test = { module = "org.jetbrains.kotlin:kotlin-test", version.ref = "kotlin" } +kotlinx-coroutines-test = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-test", version.ref = "coroutines" } +kotlinx-coroutines-debug = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-debug", version.ref = "coroutines" } +ktor-server-test-host = { group = "io.ktor", name = "ktor-server-test-host", version.ref = "ktor" } +mockk = { group = "io.mockk", name = "mockk", version.ref = "mockk" } + + + +[plugins] +kotlin-jvm = { id = "org.jetbrains.kotlin.jvm", version.ref = "kotlin" } +kotlin-serialization = { id = "org.jetbrains.kotlin.plugin.serialization", version.ref = "kotlin" } +dokka = { id = "org.jetbrains.dokka", version.ref = "dokka" } \ No newline at end of file diff --git a/samples/cli.kt b/samples/cli.kt new file mode 100644 index 0000000..b52b805 --- /dev/null +++ b/samples/cli.kt @@ -0,0 +1,297 @@ +package org.jetbrains.kotlinx.mcp + +import org.jetbrains.kotlinx.mcp.client.Client +import org.jetbrains.kotlinx.mcp.client.StdioClientTransport +import io.ktor.http.* +import io.ktor.server.application.* +import io.ktor.server.cio.* +import io.ktor.server.engine.* +import io.ktor.server.response.* +import io.ktor.server.routing.* +import io.ktor.server.sse.* +import io.ktor.util.collections.* +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.runBlocking +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.jsonObject +import org.jetbrains.kotlinx.mcp.server.MCP +import org.jetbrains.kotlinx.mcp.server.SSEServerTransport +import org.jetbrains.kotlinx.mcp.server.Server +import org.jetbrains.kotlinx.mcp.server.ServerOptions +import org.jetbrains.kotlinx.mcp.server.StdioServerTransport + +fun main(args: Array) { + if (args.isEmpty()) + return + val first = args[0] + when (first) { + "--server" -> runServer() + "--demo" -> runDemo() + "--sse-server" -> { + if (args.size < 2) { + System.err.println("Missing port argument") + return + } + val port = args[1].toIntOrNull() + if (port == null) { + System.err.println("Invalid port: ${args[1]}") + return + } + runSseServer(port) + } + + else -> { + System.err.println("Unknown argument: $first") + } + } +} + +private fun runDemo() { + val processBuilder = ProcessBuilder("npx", "-y", "@jetbrains/mcp-proxy") + + var process: Process? = null + try { + process = processBuilder.start() + + val client = Client( + Implementation("test", "1.0"), + ) + val transport = StdioClientTransport(process.inputStream, process.outputStream) + runBlocking { + client.connect(transport) + + val serverCapabilities = client.getServerCapabilities() + + // Resources capability check + serverCapabilities?.resources?.let { + try { + val resources = client.listResources() + System.err.println("Resources: ${resources?.resources?.joinToString { it.name }}") + } catch (e: Exception) { + System.err.println("Failed to list resources: ${e.message}") + } + } + + // Tools capability check + serverCapabilities?.tools?.let { + try { + val terminal = client.callTool("execute_terminal_command", mapOf("command" to "ls")) + System.err.println(terminal?.content?.first()) + + val tools = client.listTools() + System.err.println(tools?.tools?.joinToString(", ") { tool -> tool.name }) + +// tools?.tools?.reversed()?.find { it.name == "toggle_debugger_breakpoint" }?.let { org.jetbrains.kotlinx.mcp.callTool(client, it) } + + tools?.tools?.reversed()?.forEachIndexed { i, tool -> + System.err.println("$i out of ${tools.tools.size}: ${tool.name}") + callTool(client, tool) + } + } catch (e: Exception) { + System.err.println("Failed to list tools: ${e.message}") + } + } + + // Prompts capability check + serverCapabilities?.prompts?.let { + try { + val prompts = client.listPrompts() + System.err.println("Prompts: ${prompts?.prompts?.joinToString { it.name }}") + } catch (e: Exception) { + System.err.println("Failed to list prompts: ${e.message}") + } + } + } + } finally { + process?.destroy() + } +} + +private suspend fun callTool(client: Client, tool: Tool) { + System.err.println(tool.inputSchema) + + val map = fillSchema(tool.inputSchema) + + System.err.println("calling: ${tool.name}: $map") + val result = try { + client.callTool(CallToolRequest(tool.name, map)) + } catch (e: Exception) { + System.err.println("Failed to call tool ${tool.name}: ${e.message}") + return + } + System.err.println("Result: ${result?.content?.first()}\n") +} + +private fun fillSchema(schema: Tool.Input): JsonObject { + return buildJsonObject { + schema.properties.forEach { name, elt -> + val type = (elt.jsonObject["type"] as JsonPrimitive).content + val value = when (type) { + "string" -> JsonPrimitive("Hello") + "number" -> JsonPrimitive(42) + "boolean" -> JsonPrimitive(true) + else -> { + System.err.println("+".repeat(30) + " Unknown type: $type " + "+".repeat(30)) + JsonPrimitive("Unknown") + } + } + put(name, value) + } + } + +} + +private fun runServer() { + val def = CompletableDeferred() + + val server = Server( + Implementation( + name = "mcp-kotlin test server", + version = "0.1.0" + ), + ServerOptions( + capabilities = ServerCapabilities( + prompts = ServerCapabilities.Prompts(listChanged = true), + resources = ServerCapabilities.Resources(subscribe = true, listChanged = true), + tools = ServerCapabilities.Tools(listChanged = true), + ) + ), + onCloseCallback = { + def.complete(Unit) + } + ) + + server.addPrompt( + name = "Kotlin Developer", + description = "Develop small kotlin applications", + arguments = listOf( + PromptArgument( + name = "Project Name", + description = "Project name for the new project", + required = true + ) + ) + ) { request -> + GetPromptResult( + "Description for ${request.name}", + messages = listOf( + PromptMessage( + role = Role.user, + content = TextContent("Develop a kotlin project named ${request.arguments?.get("Project Name")}") + ) + ) + ) + } + + // Add a tool + server.addTool( + name = "Test org.jetbrains.kotlinx.mcp.Tool", + description = "A test tool", + inputSchema = Tool.Input() + ) { request -> + CallToolResult( + content = listOf(TextContent("Hello, world!")) + ) + } + + // Add a resource + server.addResource( + uri = "https://google.com/", + name = "Google Search", + description = "Web search engine", + mimeType = "text/html" + ) { request -> + ReadResourceResult( + contents = listOf( + TextResourceContents("Placeholder content for ${request.uri}", request.uri, "text/html") + ) + ) + } + + // Note: The server will handle listing prompts, tools, and resources automatically. + // The handleListResourceTemplates will return empty as defined in the Server code. + + val transport = StdioServerTransport() + + val err = System.err + + runBlocking { + server.connect(transport) + err.println("Server running on stdio") + def.await() + } + + err.println("Server closed") +} + +fun runSseServer(port: Int): Unit = runBlocking { + val servers = ConcurrentMap() + + embeddedServer(CIO, host = "0.0.0.0", port = port) { + install(SSE) + routing { + sse("/sse") { + val transport = SSEServerTransport("/message", this) + val options = ServerOptions( + capabilities = ServerCapabilities( + prompts = ServerCapabilities.Prompts(listChanged = null), + resources = ServerCapabilities.Resources(subscribe = null, listChanged = null), + ) + ) + val server = Server( + Implementation( + name = "mcp-kotlin test server", + version = "0.1.0" + ), + options + ) + + // For SSE, you can also add prompts/tools/resources if needed: + // server.addTool(...), server.addPrompt(...), server.addResource(...) + + servers[transport.sessionId] = server + + server.onCloseCallback = { + println("Server closed") + servers.remove(transport.sessionId) + } + + server.connect(transport) + } + post("/message") { + println("Received Message") + val sessionId: String = call.request.queryParameters["sessionId"]!! + val transport = servers[sessionId]?.transport as? SSEServerTransport + if (transport == null) { + call.respond(HttpStatusCode.NotFound, "Session not found") + return@post + } + + transport.handlePostMessage(call) + } + } + }.start(wait = true) +} + +fun runSseServerKtor(port: Int): Unit = runBlocking { + embeddedServer(CIO, host = "0.0.0.0", port = port) { + MCP { + val options = ServerOptions( + capabilities = ServerCapabilities( + prompts = ServerCapabilities.Prompts(listChanged = null), + resources = ServerCapabilities.Resources(subscribe = null, listChanged = null), + ) + ) + + Server( + Implementation( + name = "mcp-kotlin test server", + version = "0.1.0" + ), + options + ) + } + } +} \ No newline at end of file diff --git a/samples/connect-to-intellij.sh b/samples/connect-to-intellij.sh new file mode 100755 index 0000000..817b1bd --- /dev/null +++ b/samples/connect-to-intellij.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash + +../gradlew clean jar; java -jar ../build/libs/mcp-kotlin-sdk-1.0-SNAPSHOT.jar --demo \ No newline at end of file diff --git a/samples/inspector.sh b/samples/inspector.sh new file mode 100755 index 0000000..dc6de14 --- /dev/null +++ b/samples/inspector.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash + +../gradlew clean jar && npx @modelcontextprotocol/inspector java -jar ../build/libs/mcp-kotlin-sdk-1.0-SNAPSHOT.jar --server \ No newline at end of file diff --git a/samples/node-server/build.sh b/samples/node-server/build.sh new file mode 100755 index 0000000..1142cc7 --- /dev/null +++ b/samples/node-server/build.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash + +pnpm build \ No newline at end of file diff --git a/samples/node-server/package-lock.json b/samples/node-server/package-lock.json new file mode 100644 index 0000000..52af63c --- /dev/null +++ b/samples/node-server/package-lock.json @@ -0,0 +1,528 @@ +{ + "name": "node-server", + "version": "1.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "dependencies": { + "@modelcontextprotocol/sdk": "1.0.3", + "node-fetch": "^3.3.2" + }, + "bin": { + "node-server": "dist/server.js" + }, + "devDependencies": { + "@types/node": "^20.11.0", + "shx": "^0.3.4", + "typescript": "^5.3.3" + } + }, + "node_modules/@modelcontextprotocol/sdk": { + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.0.3.tgz", + "integrity": "sha512-2as3cX/VJ0YBHGmdv3GFyTpoM8q2gqE98zh3Vf1NwnsSY0h3mvoO07MUzfygCKkWsFjcZm4otIiqD6Xh7kiSBQ==", + "license": "MIT", + "dependencies": { + "content-type": "^1.0.5", + "raw-body": "^3.0.0", + "zod": "^3.23.8" + } + }, + "node_modules/@types/node": { + "version": "20.17.10", + "resolved": "https://registry.npmjs.org/@types/node/-/node-20.17.10.tgz", + "integrity": "sha512-/jrvh5h6NXhEauFFexRin69nA0uHJ5gwk4iDivp/DeoEua3uwCUto6PC86IpRITBOs4+6i2I56K5x5b6WYGXHA==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.19.2" + } + }, + "node_modules/balanced-match": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz", + "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==", + "dev": true, + "license": "MIT" + }, + "node_modules/brace-expansion": { + "version": "1.1.11", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.11.tgz", + "integrity": "sha512-iCuPHDFgrHX7H2vEI/5xpz07zSHB00TpugqhmYtVmMO6518mCuRMoOYFldEBl0g187ufozdaHgWKcYFb61qGiA==", + "dev": true, + "license": "MIT", + "dependencies": { + "balanced-match": "^1.0.0", + "concat-map": "0.0.1" + } + }, + "node_modules/bytes": { + "version": "3.1.2", + "resolved": "https://registry.npmjs.org/bytes/-/bytes-3.1.2.tgz", + "integrity": "sha512-/Nf7TyzTx6S3yRJObOAV7956r8cr2+Oj8AC5dt8wSP3BQAoeX58NoHyCU8P8zGkNXStjTSi6fzO6F0pBdcYbEg==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/concat-map": { + "version": "0.0.1", + "resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz", + "integrity": "sha512-/Srv4dswyQNBfohGpz9o6Yb3Gz3SrUDqBH5rTuhGR7ahtlbYKnVxw2bCFMRljaA7EXHaXZ8wsHdodFvbkhKmqg==", + "dev": true, + "license": "MIT" + }, + "node_modules/content-type": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-1.0.5.tgz", + "integrity": "sha512-nTjqfcBFEipKdXCv4YDQWCfmcLZKm81ldF0pAopTvyrFGVbcR6P/VAAd5G7N+0tTr8QqiU0tFadD6FK4NtJwOA==", + "license": "MIT", + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/data-uri-to-buffer": { + "version": "4.0.1", + "resolved": "https://registry.npmjs.org/data-uri-to-buffer/-/data-uri-to-buffer-4.0.1.tgz", + "integrity": "sha512-0R9ikRb668HB7QDxT1vkpuUBtqc53YyAwMwGeUFKRojY/NWKvdZ+9UYtRfGmhqNbRkTSVpMbmyhXipFFv2cb/A==", + "license": "MIT", + "engines": { + "node": ">= 12" + } + }, + "node_modules/depd": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", + "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/fetch-blob": { + "version": "3.2.0", + "resolved": "https://registry.npmjs.org/fetch-blob/-/fetch-blob-3.2.0.tgz", + "integrity": "sha512-7yAQpD2UMJzLi1Dqv7qFYnPbaPx7ZfFK6PiIxQ4PfkGPyNyl2Ugx+a/umUonmKqjhM4DnfbMvdX6otXq83soQQ==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/jimmywarting" + }, + { + "type": "paypal", + "url": "https://paypal.me/jimmywarting" + } + ], + "license": "MIT", + "dependencies": { + "node-domexception": "^1.0.0", + "web-streams-polyfill": "^3.0.3" + }, + "engines": { + "node": "^12.20 || >= 14.13" + } + }, + "node_modules/formdata-polyfill": { + "version": "4.0.10", + "resolved": "https://registry.npmjs.org/formdata-polyfill/-/formdata-polyfill-4.0.10.tgz", + "integrity": "sha512-buewHzMvYL29jdeQTVILecSaZKnt/RJWjoZCF5OW60Z67/GmSLBkOFM7qh1PI3zFNtJbaZL5eQu1vLfazOwj4g==", + "license": "MIT", + "dependencies": { + "fetch-blob": "^3.1.2" + }, + "engines": { + "node": ">=12.20.0" + } + }, + "node_modules/fs.realpath": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/fs.realpath/-/fs.realpath-1.0.0.tgz", + "integrity": "sha512-OO0pH2lK6a0hZnAdau5ItzHPI6pUlvI7jMVnxUQRtw4owF2wk8lOSabtGDCTP4Ggrg2MbGnWO9X8K1t4+fGMDw==", + "dev": true, + "license": "ISC" + }, + "node_modules/function-bind": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz", + "integrity": "sha512-7XHNxH7qX9xG5mIwxkhumTox/MIRNcOgDrxWsMt2pAr23WHp6MrRlN7FBSFpCpr+oVO0F744iUgR82nJMfG2SA==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/glob": { + "version": "7.2.3", + "resolved": "https://registry.npmjs.org/glob/-/glob-7.2.3.tgz", + "integrity": "sha512-nFR0zLpU2YCaRxwoCJvL6UvCH2JFyFVIvwTLsIf21AuHlMskA1hhTdk+LlYJtOlYt9v6dvszD2BGRqBL+iQK9Q==", + "deprecated": "Glob versions prior to v9 are no longer supported", + "dev": true, + "license": "ISC", + "dependencies": { + "fs.realpath": "^1.0.0", + "inflight": "^1.0.4", + "inherits": "2", + "minimatch": "^3.1.1", + "once": "^1.3.0", + "path-is-absolute": "^1.0.0" + }, + "engines": { + "node": "*" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" + } + }, + "node_modules/hasown": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.2.tgz", + "integrity": "sha512-0hJU9SCPvmMzIBdZFqNPXWa6dqh7WdH0cII9y+CyS8rG3nL48Bclra9HmKhVVUHyPWNH5Y7xDwAB7bfgSjkUMQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "function-bind": "^1.1.2" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/http-errors": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.0.tgz", + "integrity": "sha512-FtwrG/euBzaEjYeRqOgly7G0qviiXoJWnvEH2Z1plBdXgbyjv34pHTSb9zoeHMyDy33+DWy5Wt9Wo+TURtOYSQ==", + "license": "MIT", + "dependencies": { + "depd": "2.0.0", + "inherits": "2.0.4", + "setprototypeof": "1.2.0", + "statuses": "2.0.1", + "toidentifier": "1.0.1" + }, + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/iconv-lite": { + "version": "0.6.3", + "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.6.3.tgz", + "integrity": "sha512-4fCk79wshMdzMp2rH06qWrJE4iolqLhCUH+OiuIgU++RB0+94NlDL81atO7GX55uUKueo0txHNtvEyI6D7WdMw==", + "license": "MIT", + "dependencies": { + "safer-buffer": ">= 2.1.2 < 3.0.0" + }, + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/inflight": { + "version": "1.0.6", + "resolved": "https://registry.npmjs.org/inflight/-/inflight-1.0.6.tgz", + "integrity": "sha512-k92I/b08q4wvFscXCLvqfsHCrjrF7yiXsQuIVvVE7N82W3+aqpzuUdBbfhWcy/FZR3/4IgflMgKLOsvPDrGCJA==", + "deprecated": "This module is not supported, and leaks memory. Do not use it. Check out lru-cache if you want a good and tested way to coalesce async requests by a key value, which is much more comprehensive and powerful.", + "dev": true, + "license": "ISC", + "dependencies": { + "once": "^1.3.0", + "wrappy": "1" + } + }, + "node_modules/inherits": { + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", + "license": "ISC" + }, + "node_modules/interpret": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/interpret/-/interpret-1.4.0.tgz", + "integrity": "sha512-agE4QfB2Lkp9uICn7BAqoscw4SZP9kTE2hxiFI3jBPmXJfdqiahTbUuKGsMoN2GtqL9AxhYioAcVvgsb1HvRbA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.10" + } + }, + "node_modules/is-core-module": { + "version": "2.15.1", + "resolved": "https://registry.npmjs.org/is-core-module/-/is-core-module-2.15.1.tgz", + "integrity": "sha512-z0vtXSwucUJtANQWldhbtbt7BnL0vxiFjIdDLAatwhDYty2bad6s+rijD6Ri4YuYJubLzIJLUidCh09e1djEVQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "hasown": "^2.0.2" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/minimatch": { + "version": "3.1.2", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", + "integrity": "sha512-J7p63hRiAjw1NDEww1W7i37+ByIrOWO5XQQAzZ3VOcL0PNybwpfmV/N05zFAzwQ9USyEcX6t3UO+K5aqBQOIHw==", + "dev": true, + "license": "ISC", + "dependencies": { + "brace-expansion": "^1.1.7" + }, + "engines": { + "node": "*" + } + }, + "node_modules/minimist": { + "version": "1.2.8", + "resolved": "https://registry.npmjs.org/minimist/-/minimist-1.2.8.tgz", + "integrity": "sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/node-domexception": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/node-domexception/-/node-domexception-1.0.0.tgz", + "integrity": "sha512-/jKZoMpw0F8GRwl4/eLROPA3cfcXtLApP0QzLmUT/HuPCZWyB7IY9ZrMeKw2O/nFIqPQB3PVM9aYm0F312AXDQ==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/jimmywarting" + }, + { + "type": "github", + "url": "https://paypal.me/jimmywarting" + } + ], + "license": "MIT", + "engines": { + "node": ">=10.5.0" + } + }, + "node_modules/node-fetch": { + "version": "3.3.2", + "resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-3.3.2.tgz", + "integrity": "sha512-dRB78srN/l6gqWulah9SrxeYnxeddIG30+GOqK/9OlLVyLg3HPnr6SqOWTWOXKRwC2eGYCkZ59NNuSgvSrpgOA==", + "license": "MIT", + "dependencies": { + "data-uri-to-buffer": "^4.0.0", + "fetch-blob": "^3.1.4", + "formdata-polyfill": "^4.0.10" + }, + "engines": { + "node": "^12.20.0 || ^14.13.1 || >=16.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/node-fetch" + } + }, + "node_modules/once": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", + "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==", + "dev": true, + "license": "ISC", + "dependencies": { + "wrappy": "1" + } + }, + "node_modules/path-is-absolute": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/path-is-absolute/-/path-is-absolute-1.0.1.tgz", + "integrity": "sha512-AVbw3UJ2e9bq64vSaS9Am0fje1Pa8pbGqTTsmXfaIiMpnr5DlDhfJOuLj9Sf95ZPVDAUerDfEk88MPmPe7UCQg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/path-parse": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/path-parse/-/path-parse-1.0.7.tgz", + "integrity": "sha512-LDJzPVEEEPR+y48z93A0Ed0yXb8pAByGWo/k5YYdYgpY2/2EsOsksJrq7lOHxryrVOn1ejG6oAp8ahvOIQD8sw==", + "dev": true, + "license": "MIT" + }, + "node_modules/raw-body": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-3.0.0.tgz", + "integrity": "sha512-RmkhL8CAyCRPXCE28MMH0z2PNWQBNk2Q09ZdxM9IOOXwxwZbN+qbWaatPkdkWIKL2ZVDImrN/pK5HTRz2PcS4g==", + "license": "MIT", + "dependencies": { + "bytes": "3.1.2", + "http-errors": "2.0.0", + "iconv-lite": "0.6.3", + "unpipe": "1.0.0" + }, + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/rechoir": { + "version": "0.6.2", + "resolved": "https://registry.npmjs.org/rechoir/-/rechoir-0.6.2.tgz", + "integrity": "sha512-HFM8rkZ+i3zrV+4LQjwQ0W+ez98pApMGM3HUrN04j3CqzPOzl9nmP15Y8YXNm8QHGv/eacOVEjqhmWpkRV0NAw==", + "dev": true, + "dependencies": { + "resolve": "^1.1.6" + }, + "engines": { + "node": ">= 0.10" + } + }, + "node_modules/resolve": { + "version": "1.22.8", + "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.22.8.tgz", + "integrity": "sha512-oKWePCxqpd6FlLvGV1VU0x7bkPmmCNolxzjMf4NczoDnQcIWrAF+cPtZn5i6n+RfD2d9i0tzpKnG6Yk168yIyw==", + "dev": true, + "license": "MIT", + "dependencies": { + "is-core-module": "^2.13.0", + "path-parse": "^1.0.7", + "supports-preserve-symlinks-flag": "^1.0.0" + }, + "bin": { + "resolve": "bin/resolve" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/safer-buffer": { + "version": "2.1.2", + "resolved": "https://registry.npmjs.org/safer-buffer/-/safer-buffer-2.1.2.tgz", + "integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==", + "license": "MIT" + }, + "node_modules/setprototypeof": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/setprototypeof/-/setprototypeof-1.2.0.tgz", + "integrity": "sha512-E5LDX7Wrp85Kil5bhZv46j8jOeboKq5JMmYM3gVGdGH8xFpPWXUMsNrlODCrkoxMEeNi/XZIwuRvY4XNwYMJpw==", + "license": "ISC" + }, + "node_modules/shelljs": { + "version": "0.8.5", + "resolved": "https://registry.npmjs.org/shelljs/-/shelljs-0.8.5.tgz", + "integrity": "sha512-TiwcRcrkhHvbrZbnRcFYMLl30Dfov3HKqzp5tO5b4pt6G/SezKcYhmDg15zXVBswHmctSAQKznqNW2LO5tTDow==", + "dev": true, + "license": "BSD-3-Clause", + "dependencies": { + "glob": "^7.0.0", + "interpret": "^1.0.0", + "rechoir": "^0.6.2" + }, + "bin": { + "shjs": "bin/shjs" + }, + "engines": { + "node": ">=4" + } + }, + "node_modules/shx": { + "version": "0.3.4", + "resolved": "https://registry.npmjs.org/shx/-/shx-0.3.4.tgz", + "integrity": "sha512-N6A9MLVqjxZYcVn8hLmtneQWIJtp8IKzMP4eMnx+nqkvXoqinUPCbUFLp2UcWTEIUONhlk0ewxr/jaVGlc+J+g==", + "dev": true, + "license": "MIT", + "dependencies": { + "minimist": "^1.2.3", + "shelljs": "^0.8.5" + }, + "bin": { + "shx": "lib/cli.js" + }, + "engines": { + "node": ">=6" + } + }, + "node_modules/statuses": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", + "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/supports-preserve-symlinks-flag": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/supports-preserve-symlinks-flag/-/supports-preserve-symlinks-flag-1.0.0.tgz", + "integrity": "sha512-ot0WnXS9fgdkgIcePe6RHNk1WA8+muPa6cSjeR3V8K27q9BB1rTE3R1p7Hv0z1ZyAc8s6Vvv8DIyWf681MAt0w==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/toidentifier": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/toidentifier/-/toidentifier-1.0.1.tgz", + "integrity": "sha512-o5sSPKEkg/DIQNmH43V0/uerLrpzVedkUh8tGNvaeXpfpuwjKenlSox/2O/BTlZUtEe+JG7s5YhEz608PlAHRA==", + "license": "MIT", + "engines": { + "node": ">=0.6" + } + }, + "node_modules/typescript": { + "version": "5.7.2", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.7.2.tgz", + "integrity": "sha512-i5t66RHxDvVN40HfDd1PsEThGNnlMCMT3jMUuoh9/0TaqWevNontacunWyN02LA9/fIbEWlcHZcgTKb9QoaLfg==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.19.8", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.19.8.tgz", + "integrity": "sha512-ve2KP6f/JnbPBFyobGHuerC9g1FYGn/F8n1LWTwNxCEzd6IfqTwUQcNXgEtmmQ6DlRrC1hrSrBnCZPokRrDHjw==", + "dev": true, + "license": "MIT" + }, + "node_modules/unpipe": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/unpipe/-/unpipe-1.0.0.tgz", + "integrity": "sha512-pjy2bYhSsufwWlKwPc+l3cN7+wuJlK6uz0YdJEOlQDbl6jo/YlPi4mb8agUkVC8BF7V8NuzeyPNqRksA3hztKQ==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/web-streams-polyfill": { + "version": "3.3.3", + "resolved": "https://registry.npmjs.org/web-streams-polyfill/-/web-streams-polyfill-3.3.3.tgz", + "integrity": "sha512-d2JWLCivmZYTSIoge9MsgFCZrt571BikcWGYkjC1khllbTeDlGqZ2D8vD8E/lJa8WGWbb7Plm8/XJYV7IJHZZw==", + "license": "MIT", + "engines": { + "node": ">= 8" + } + }, + "node_modules/wrappy": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz", + "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", + "dev": true, + "license": "ISC" + }, + "node_modules/zod": { + "version": "3.24.1", + "resolved": "https://registry.npmjs.org/zod/-/zod-3.24.1.tgz", + "integrity": "sha512-muH7gBL9sI1nciMZV67X5fTKKBLtwpZ5VBp1vsOQzj1MhrBZ4wlVCm3gedKZWLp0Oyel8sIGfeiz54Su+OVT+A==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/colinhacks" + } + } + } +} diff --git a/samples/node-server/package.json b/samples/node-server/package.json new file mode 100644 index 0000000..2174916 --- /dev/null +++ b/samples/node-server/package.json @@ -0,0 +1,26 @@ +{ + "main": "dist/server.js", + "type": "module", + "bin": { + "node-server": "dist/server.js" + }, + "files": [ + "dist" + ], + "scripts": { + "build": "tsc --outDir dist && shx chmod +x dist/*.js", + "prepare": "npm run build", + "watch": "tsc --watch --outDir dist", + "run_server": "dist/server.js", + "run_client": "dist/client.js" + }, + "dependencies": { + "@modelcontextprotocol/sdk": "1.0.3", + "node-fetch": "^3.3.2" + }, + "devDependencies": { + "@types/node": "^20.11.0", + "shx": "^0.3.4", + "typescript": "^5.3.3" + } +} diff --git a/samples/node-server/run_client.sh b/samples/node-server/run_client.sh new file mode 100755 index 0000000..528de14 --- /dev/null +++ b/samples/node-server/run_client.sh @@ -0,0 +1,2 @@ +#!/bin/bash +./dist/client.js \ No newline at end of file diff --git a/samples/node-server/run_server.sh b/samples/node-server/run_server.sh new file mode 100755 index 0000000..428590d --- /dev/null +++ b/samples/node-server/run_server.sh @@ -0,0 +1,2 @@ +#!/bin/bash +./dist/server.js \ No newline at end of file diff --git a/samples/node-server/src/client.ts b/samples/node-server/src/client.ts new file mode 100644 index 0000000..0a7cd1b --- /dev/null +++ b/samples/node-server/src/client.ts @@ -0,0 +1,52 @@ +#!/usr/bin/env node + +import * as net from 'net'; +import { Readable, Writable } from 'stream'; + +const client = new net.Socket(); +const PORT = 3000; +const HOST = '127.0.0.1'; + +// If startClient is supposed to be asynchronous and do something later, +// keep async. Otherwise, remove async if not needed. +function startClient(input: Readable, output: Writable) { + // TODO: Implement client logic here if needed + // For now, it's just a placeholder. +} + +const input = new Readable({ + read() { + // No-op: We'll push data into this stream externally. + } +}); + +const output = new Writable({ + write(chunk, encoding, callback) { + console.log('Writing to output (and pushing to input):', chunk.toString('utf-8')); + input.push(chunk); + callback(); + } +}); + +client.connect(PORT, HOST, () => { + console.log('Connected to server'); + startClient(input, output); +}); + +client.on('data', (data) => { + // When we receive data from the server socket, push it into `input`. + input.push(data); +}); + +client.on('close', () => { + input.push(null); // Signal the end of the input stream + console.log('Connection closed by server'); +}); + +client.on('error', (err) => { + console.error('Client error:', err); +}); + +// Push initial data into `input`. If no reader is attached, this might not be visible. +// Consider adding `input.on('data', ...)` somewhere else if you want to see this data. +input.push("Hello"); \ No newline at end of file diff --git a/samples/node-server/src/server.ts b/samples/node-server/src/server.ts new file mode 100755 index 0000000..7aec7d9 --- /dev/null +++ b/samples/node-server/src/server.ts @@ -0,0 +1,95 @@ +#!/usr/bin/env node + +import * as net from 'net'; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import { Readable, Writable, WritableOptions } from 'stream'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js'; + +const PORT = 3000; + +const mcpServer = new Server( + { + name: "jetbrains/proxy", + version: "0.1.0", + }, + { + capabilities: { + tools: { + listChanged: true, + }, + resources: {}, + }, + }, +); + +mcpServer.setRequestHandler(ListToolsRequestSchema, async () => { + console.log("ListToolsRequest received"); + return { + tools: [] + }; +}); + +async function handleConnection(stdin: Readable, stdout: Writable) { + const transport = new StdioServerTransport(stdin, stdout); + try { + await mcpServer.connect(transport); + } catch (err) { + console.error('Error connecting MCP server:', err); + } +} + +// A Writable stream class that writes data back to the socket +class SocketWritable extends Writable { + private socket: net.Socket; + + constructor(socket: net.Socket, options?: WritableOptions) { + super(options); + this.socket = socket; + } + + _write(chunk: any, encoding: BufferEncoding, callback: (error?: Error | null) => void) { + console.log("Socket write:", chunk.toString()); + this.socket.write(chunk, encoding, callback); + } +} + +const server = net.createServer((socket) => { + console.log('Client connected'); + + // Create a Readable stream that we will push data into from the socket + const stdin = new Readable({ + read() { + // We'll use stdin.push() from socket data events + }, + }); + + // Create a Writable that writes back to the socket + const stdout = new SocketWritable(socket); + + socket.on('data', (chunk) => { + const decoded = chunk.toString("utf-8"); + console.log("Socket data:", decoded); + // Push the raw buffer into stdin so the server can read it + stdin.push(chunk); + }); + + socket.on('end', () => { + stdin.push(null); // Signal EOF to the stdin stream + console.log('Client disconnected'); + }); + + socket.on('error', (err) => { + console.error('Socket error:', err); + }); + + handleConnection(stdin, stdout); +}); + +server.listen(PORT, () => { + console.log(`MCP proxy server listening on port ${PORT}`); +}); + +server.on('error', (err) => { + console.error('Server error:', err); +}); \ No newline at end of file diff --git a/samples/node-server/src/tmp.ts b/samples/node-server/src/tmp.ts new file mode 100644 index 0000000..a36d6b7 --- /dev/null +++ b/samples/node-server/src/tmp.ts @@ -0,0 +1,151 @@ +#!/usr/bin/env node +import { Server } from "@modelcontextprotocol/sdk/server/index.js"; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import { + CallToolRequestSchema, + CallToolResult, + ListToolsRequestSchema, +} from "@modelcontextprotocol/sdk/types.js"; + +interface IDEResponseOk { + status: string; + error: null; +} + +interface IDEResponseErr { + status: null; + error: string; +} + +type IDEResponse = IDEResponseOk | IDEResponseErr; + +/** + * Try to find a working IDE endpoint. + * Logic: + * 1. If process.env.IDE_PORT is set, use that port directly. + * 2. If not set, try ports from 63342 to 63352. + * 3. For each port, send a test request to /mcp/list_tools. If it works (res.ok), use that port. + * 4. If no port works, throw an error. + */ +async function findWorkingIDEEndpoint(): Promise { + // If user specified a port, just use that + if (process.env.IDE_PORT) { + const testEndpoint = `http://localhost:${process.env.IDE_PORT}/api`; + if (await testListTools(testEndpoint)) { + return testEndpoint; + } else { + throw new Error(`Specified IDE_PORT=${process.env.IDE_PORT} but it is not responding correctly.`); + } + } + + for (let port = 63342; port <= 63352; port++) { + const candidateEndpoint = `http://localhost:${port}/api`; + if (await testListTools(candidateEndpoint)) { + return candidateEndpoint; + } + } + sendToolsChanged() + previousResponse = "" + throw new Error("No working IDE endpoint found in range 63342-63352"); +} + +let previousResponse: string | null = null; + +function sendToolsChanged() { + try { + server.notification({ method: "notifications/tools/list_changed" }) + } catch (_) { + } +} + +async function testListTools(endpoint: string): Promise { + try { + const res = await fetch(`${endpoint}/mcp/list_tools`); + + if (!res.ok) { + return false; + } + + const currentResponse = await res.text(); + if (previousResponse !== null && previousResponse !== currentResponse) { + sendToolsChanged() + } + previousResponse = currentResponse; + return true; + } catch { + return false; + } +} + + +const server = new Server( + { + name: "jetbrains/proxy", + version: "0.1.0", + }, + { + capabilities: { + tools: { + listChanged: true, + }, + resources: {}, + }, + }, +); + +server.setRequestHandler(ListToolsRequestSchema, async () => { + const endpoint = await findWorkingIDEEndpoint(); + return { + tools: await fetch(`${endpoint}/mcp/list_tools`) + .then(res => res.ok ? res.json() : Promise.reject(new Error("Unable to list tools"))) + } +}); + +async function handleToolCall(name: string, args: any): Promise { + try { + const endPoint = await findWorkingIDEEndpoint(); + const response = await fetch(`${endPoint}/mcp/${name}`, { + method: 'POST', + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(args) + }); + + if (!response.ok) { + throw new Error(`Response failed: ${response.status}`); + } + + const { status, error }: IDEResponse = await response.json(); + const isError = !!error; + const text = status ?? error; + return { + content: [{ type: "text", text: text }], + isError, + }; + } catch (error: any) { + return { + content: [{ + type: "text", + text: error instanceof Error ? error.message : "Unknown error", + }], + isError: true, + }; + } +} + +server.setRequestHandler(CallToolRequestSchema, async (request) => + handleToolCall(request.params.name, request.params.arguments ?? {}) +); + +async function runServer() { + const transport = new StdioServerTransport(); + await server.connect(transport); + const checkEndpoint = () => findWorkingIDEEndpoint().catch(); + // We need to recheck the IDE endpoint every 10 seconds since IDE might be closed or restarted + setInterval(checkEndpoint, 10000); + await checkEndpoint(); + console.error("JetBrains Proxy MCP Server running on stdio"); +} + +runServer().catch(console.error); \ No newline at end of file diff --git a/samples/node-server/tsconfig.json b/samples/node-server/tsconfig.json new file mode 100644 index 0000000..feb5dc6 --- /dev/null +++ b/samples/node-server/tsconfig.json @@ -0,0 +1,19 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "Node16", + "moduleResolution": "Node16", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "resolveJsonModule": true, + "noImplicitAny": false + }, + "include": [ + "src/**/*" + ], + "exclude": [ + "node_modules" + ] +} \ No newline at end of file diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/Client.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/Client.kt new file mode 100644 index 0000000..4ef15b7 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/Client.kt @@ -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(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(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(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(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(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(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(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(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(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(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(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(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, + 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(request, options) + } else { + request(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(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()) + } +} \ No newline at end of file diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/SSEClientTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/SSEClientTransport.kt new file mode 100644 index 0000000..8114940 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/SSEClientTransport.kt @@ -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() + + 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(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() + } +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/StdioClientTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/StdioClientTransport.kt new file mode 100644 index 0000000..2c1328f --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/StdioClientTransport.kt @@ -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(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." } + } + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketClientTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketClientTransport.kt new file mode 100644 index 0000000..1db6a6e --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketClientTransport.kt @@ -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) + } + } +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketMcpKtorClientExtensions.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketMcpKtorClientExtensions.kt new file mode 100644 index 0000000..0b0a7db --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/WebSocketMcpKtorClientExtensions.kt @@ -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 +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/sse.ktor.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/sse.ktor.kt new file mode 100644 index 0000000..f17af26 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/client/sse.ktor.kt @@ -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 +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/McpKtorServerPlugin.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/McpKtorServerPlugin.kt new file mode 100644 index 0000000..713f7e1 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/McpKtorServerPlugin.kt @@ -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() + + 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" } + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/SSEServerTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/SSEServerTransport.kt new file mode 100644 index 0000000..4412368 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/SSEServerTransport.kt @@ -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(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), + ) + } +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/Server.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/Server.kt new file mode 100644 index 0000000..e709297 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/Server.kt @@ -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(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() + private val prompts = mutableMapOf() + private val resources = mutableMapOf() + + /** + * 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(Method.Defined.Initialize) { request, _ -> + handleInitialize(request) + } + setNotificationHandler(Method.Defined.NotificationsInitialized) { + onInitialized?.invoke() + CompletableDeferred(Unit) + } + + // Internal handlers for tools + if (capabilities.tools != null) { + setRequestHandler(Method.Defined.ToolsList) { _, _ -> + handleListTools() + } + setRequestHandler(Method.Defined.ToolsCall) { request, _ -> + handleCallTool(request) + } + } + + // Internal handlers for prompts + if (capabilities.prompts != null) { + setRequestHandler(Method.Defined.PromptsList) { _, _ -> + handleListPrompts() + } + setRequestHandler(Method.Defined.PromptsGet) { request, _ -> + handleGetPrompt(request) + } + } + + // Internal handlers for resources + if (capabilities.resources != null) { + setRequestHandler(Method.Defined.ResourcesList) { _, _ -> + handleListResources() + } + setRequestHandler(Method.Defined.ResourcesRead) { request, _ -> + handleReadResource(request) + } + setRequestHandler(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) { + 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? = 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) { + 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) { + 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(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(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(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 +) \ No newline at end of file diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/StdioServerTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/StdioServerTransport.kt new file mode 100644 index 0000000..83de52a --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/StdioServerTransport.kt @@ -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(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() + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpKtorServerExtensions.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpKtorServerExtensions.kt new file mode 100644 index 0000000..0188a80 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpKtorServerExtensions.kt @@ -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) +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpServerTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpServerTransport.kt new file mode 100644 index 0000000..c0efc1c --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/server/WebSocketMcpServerTransport.kt @@ -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") + } + } +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Protocol.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Protocol.kt new file mode 100644 index 0000000..d3be22f --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Protocol.kt @@ -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( + @PublishedApi internal val _options: ProtocolOptions? +) { + var transport: Transport? = null + private set + + @PublishedApi + internal val requestHandlers: MutableMap SendResultT?> = + mutableMapOf() + val notificationHandlers: MutableMap Unit> = + mutableMapOf() + + @PublishedApi + internal val responseHandlers: MutableMap Unit> = + mutableMapOf() + + @PublishedApi + internal val progressHandlers: MutableMap = 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(Method.Defined.NotificationsProgress) { notification -> + this.onProgress(notification) + COMPLETED + } + + setRequestHandler(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 request( + request: SendRequestT, + options: RequestOptions? = null, + ): T { + LOGGER.trace { "Sending request: ${request.method}" } + val result = CompletableDeferred() + 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) 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 setRequestHandler( + method: Method, + noinline block: suspend (T, RequestHandlerExtra) -> SendResultT? + ) { + assertRequestHandlerCapability(method) + + val type = typeOf() + 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 setNotificationHandler(method: Method, handler: (notification: T) -> Deferred) { + 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) + } +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/ReadBuffer.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/ReadBuffer.kt new file mode 100644 index 0000000..5867f08 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/ReadBuffer.kt @@ -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(line) +} + +internal fun serializeMessage(message: JSONRPCMessage): String { + return McpJson.encodeToString(message) + "\n" +} + diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Transport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Transport.kt new file mode 100644 index 0000000..08f3e22 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/Transport.kt @@ -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))? +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/WebSocketMcpTransport.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/WebSocketMcpTransport.kt new file mode 100644 index 0000000..ff89ba2 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/shared/WebSocketMcpTransport.kt @@ -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(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() + } +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/types.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/types.kt new file mode 100644 index 0000000..3832949 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/types.kt @@ -0,0 +1,1448 @@ +@file:Suppress("unused", "EnumEntryName") + +package org.jetbrains.kotlinx.mcp + +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.decodeFromJsonElement +import kotlinx.serialization.json.encodeToJsonElement +import org.jetbrains.kotlinx.mcp.shared.McpJson +import java.util.concurrent.atomic.AtomicLong + +const val LATEST_PROTOCOL_VERSION = "2024-11-05" + +val SUPPORTED_PROTOCOL_VERSIONS = arrayOf( + LATEST_PROTOCOL_VERSION, + "2024-10-07", +) + +const val JSONRPC_VERSION: String = "2.0" + +private val REQUEST_MESSAGE_ID = AtomicLong(0L) + +/** + * A progress token, used to associate progress notifications with the original request. + * Stores message ID. + */ +typealias ProgressToken = Long + +/** + * An opaque token used to represent a cursor for pagination. + */ +typealias Cursor = String + +/** + * Represents an entity that includes additional metadata in its responses. + */ +@Serializable +sealed interface WithMeta { + /** + * The protocol reserves this result property + * to allow clients and servers to attach additional metadata to their responses. + */ + @Suppress("PropertyName") + val _meta: JsonObject + + companion object { + val Empty = CustomMeta() + } +} + +/** + * An implementation of [WithMeta] containing custom metadata. + * + * @param _meta The JSON object holding metadata. Defaults to an empty JSON object. + */ +@Serializable +class CustomMeta(override val _meta: JsonObject = EmptyJsonObject) : WithMeta + +/** + * Represents a method in the protocol, which can be predefined or custom. + */ +@Serializable(with = RequestMethodSerializer::class) +sealed interface Method { + val value: String + + /** + * Enum of predefined methods supported by the protocol. + */ + @Serializable + enum class Defined(override val value: String) : Method { + Initialize("initialize"), + Ping("ping"), + ResourcesList("resources/list"), + ResourcesTemplatesList("resources/templates/list"), + ResourcesRead("resources/read"), + ResourcesSubscribe("resources/subscribe"), + ResourcesUnsubscribe("resources/unsubscribe"), + PromptsList("prompts/list"), + PromptsGet("prompts/get"), + NotificationsCancelled("notifications/cancelled"), + NotificationsInitialized("notifications/initialized"), + NotificationsProgress("notifications/progress"), + NotificationsMessage("notifications/message"), + NotificationsResourcesUpdated("notifications/resources/updated"), + NotificationsResourcesListChanged("notifications/resources/list_changed"), + NotificationsToolsListChanged("notifications/tools/list_changed"), + NotificationsRootsListChanged("notifications/roots/list_changed"), + NotificationsPromptsListChanged("notifications/prompts/list_changed"), + ToolsList("tools/list"), + ToolsCall("tools/call"), + LoggingSetLevel("logging/setLevel"), + SamplingCreateMessage("sampling/createMessage"), + CompletionComplete("completion/complete"), + RootsList("roots/list") + } + + /** + * Represents a custom method defined by the user. + */ + @Serializable + data class Custom(override val value: String) : Method +} + +/** + * Represents a request in the protocol. + */ +@Serializable(with = RequestPolymorphicSerializer::class) +sealed interface Request { + val method: Method +} + +/** + * Converts the request to a JSON-RPC request. + * + * @return The JSON-RPC request representation. + */ +fun Request.toJSON(): JSONRPCRequest { + return JSONRPCRequest( + method = method.value, + params = McpJson.encodeToJsonElement(this), + jsonrpc = JSONRPC_VERSION, + ) +} + +/** + * Decodes a JSON-RPC request into a protocol-specific [Request]. + * + * @return The decoded [Request] or null + */ +fun JSONRPCRequest.fromJSON(): Request? { + val serializer = selectRequestDeserializer(method) + val params = params ?: return null + return McpJson.decodeFromJsonElement(serializer, params) +} + +/** + * A custom request with a specified method. + * + * @param method The method associated with the request. + */ +@Serializable +open class CustomRequest(override val method: Method) : Request + +/** + * Represents a notification in the protocol. + */ +@Serializable(with = NotificationPolymorphicSerializer::class) +sealed interface Notification { + val method: Method +} + +/** + * Converts the notification to a JSON-RPC notification. + * + * @return The JSON-RPC notification representation. + */ +fun Notification.toJSON(): JSONRPCNotification { + val encoded = McpJson.encodeToJsonElement(this) + return JSONRPCNotification( + method.value, + params = encoded + ) +} + +/** + * Decodes a JSON-RPC notification into a protocol-specific [Notification]. + * + * @return The decoded [Notification]. + */ +fun JSONRPCNotification.fromJSON(): Notification { + return McpJson.decodeFromJsonElement(params) +} + +/** + * Represents the result of a request, including additional metadata. + */ +@Serializable(with = RequestResultPolymorphicSerializer::class) +sealed interface RequestResult : WithMeta + +/** + * An empty result for a request, containing optional metadata. + * + * @param _meta Additional metadata for the response. Defaults to an empty JSON object. + */ +@Serializable +data class EmptyRequestResult( + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult, ClientResult + +/** + * A uniquely identifying ID for a request in JSON-RPC. + */ +typealias RequestId = Long + +/** + * Represents a JSON-RPC message in the protocol. + */ +@Serializable(with = JSONRPCMessagePolymorphicSerializer::class) +sealed interface JSONRPCMessage + +/** + * A request that expects a response. + */ +@Serializable +data class JSONRPCRequest( + val id: RequestId = REQUEST_MESSAGE_ID.incrementAndGet(), + val method: String, + val params: JsonElement? = null, + val jsonrpc: String = JSONRPC_VERSION, +) : JSONRPCMessage + +/** + * A notification which does not expect a response. + */ +@Serializable +data class JSONRPCNotification( + val method: String, + val params: JsonElement = EmptyJsonObject, + val jsonrpc: String = JSONRPC_VERSION, +) : JSONRPCMessage + +/** + * A successful (non-error) response to a request. + */ +@Serializable +class JSONRPCResponse( + val id: RequestId, + val jsonrpc: String = JSONRPC_VERSION, + val result: RequestResult? = null, + val error: JSONRPCError? = null, +) : JSONRPCMessage + +/** + * An incomplete set of error codes that may appear in JSON-RPC responses. + */ +@Serializable(with = ErrorCodeSerializer::class) +sealed interface ErrorCode { + val code: Int + + @Serializable + enum class Defined(override val code: Int) : ErrorCode { + // SDK error codes + ConnectionClosed(-1), + RequestTimeout(-2), + + // Standard JSON-RPC error codes + ParseError(-32700), + InvalidRequest(-32600), + MethodNotFound(-32601), + InvalidParams(-32602), + InternalError(-32603), + ; + } + + @Serializable + data class Unknown(override val code: Int) : ErrorCode +} + +/** + * A response to a request that indicates an error occurred. + */ +@Serializable +data class JSONRPCError( + val code: ErrorCode, + val message: String, + val data: JsonObject = EmptyJsonObject, +) : JSONRPCMessage + +/* Cancellation */ +/** + * This notification can be sent by either side to indicate that it is cancelling a previously issued request. + * + * The request SHOULD still be in-flight, but due to communication latency, it is always possible that this notification MAY arrive after the request has already finished. + * + * This notification indicates that the result will be unused, so any associated processing SHOULD cease. + * + * A client MUST NOT attempt to cancel its `initialize` request. + */ +@Serializable +data class CancelledNotification( + /** + * The ID of the request to cancel. + * + * It MUST correspond to the ID of a request previously issued in the same direction. + */ + val requestId: RequestId, + /** + * An optional string describing the reason for the cancellation. This MAY be logged or presented to the user. + */ + val reason: String?, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientNotification, ServerNotification, WithMeta { + override val method: Method = Method.Defined.NotificationsCancelled +} + +/* Initialization */ +/** + * Describes the name and version of an MCP implementation. + */ +@Serializable +data class Implementation( + val name: String, + val version: String, +) + +/** + * Capabilities a client may support. + * Known capabilities are defined here, in this, but this is not a closed set: + * any client can define its own, additional capabilities. + */ +@Serializable +data class ClientCapabilities( + /** + * Experimental, non-standard capabilities that the client supports. + */ + val experimental: JsonObject? = EmptyJsonObject, + /** + * Present if the client supports sampling from an LLM. + */ + val sampling: JsonObject? = EmptyJsonObject, + /** + * Present if the client supports listing roots. + */ + val roots: Roots? = null, +) { + @Serializable + data class Roots( + /** + * Whether the client supports issuing notifications for changes to the root list. + */ + val listChanged: Boolean?, + ) +} + +/** + * Represents a request sent by the client. + */ +//@Serializable(with = ClientRequestPolymorphicSerializer::class) +interface ClientRequest : Request + +/** + * Represents a notification sent by the client. + */ +@Serializable(with = ClientNotificationPolymorphicSerializer::class) +sealed interface ClientNotification : Notification + +/** + * Represents a result returned to the client. + */ +@Serializable(with = ClientResultPolymorphicSerializer::class) +sealed interface ClientResult : RequestResult + +/** + * Represents a request sent by the server. + */ +//@Serializable(with = ServerRequestPolymorphicSerializer::class) +sealed interface ServerRequest : Request + +/** + * Represents a notification sent by the server. + */ +@Serializable(with = ServerNotificationPolymorphicSerializer::class) +sealed interface ServerNotification : Notification + +/** + * Represents a result returned by the server. + */ +@Serializable(with = ServerResultPolymorphicSerializer::class) +sealed interface ServerResult : RequestResult + +/** + * Represents a request or notification for an unknown method. + * + * @param method The method that is unknown. + */ +@Serializable +data class UnknownMethodRequestOrNotification( + override val method: Method, +) : ClientNotification, ClientRequest, ServerNotification, ServerRequest + +/** + * This request is sent from the client to the server when it first connects, asking it to begin initialization. + */ +@Serializable +data class InitializeRequest( + + val protocolVersion: String, + val capabilities: ClientCapabilities, + val clientInfo: Implementation, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.Initialize +} + +/** + * Represents the capabilities that a server can support. + * + * @property experimental Experimental, non-standard capabilities that the server supports. + * @property sampling Present if the client supports sampling from an LLM. + * @property logging Present if the server supports sending log messages to the client. + * @property prompts Capabilities related to prompt templates offered by the server. + * @property resources Capabilities related to resources available on the server. + * @property tools Capabilities related to tools that can be called on the server. + */ +@Serializable +data class ServerCapabilities( + val experimental: JsonObject? = EmptyJsonObject, + val sampling: JsonObject? = EmptyJsonObject, + val logging: JsonObject? = EmptyJsonObject, + val prompts: Prompts? = null, + val resources: Resources? = null, + val tools: Tools? = null, +) { + /** + * Capabilities related to prompt templates. + * + * @property listChanged Indicates if the server supports notifications when the prompt list changes. + */ + @Serializable + data class Prompts( + /** + * Whether this server supports issuing notifications for changes to the prompt list. + */ + val listChanged: Boolean?, + ) + + /** + * Capabilities related to resources. + * + * @property subscribe Indicates if clients can subscribe to resource updates. + * @property listChanged Indicates if the server supports notifications when the resource list changes. + */ + @Serializable + data class Resources( + /** + * Whether this server supports clients subscribing to resource updates. + */ + val subscribe: Boolean?, + /** + * Whether this server supports issuing notifications for changes to the resource list. + */ + val listChanged: Boolean?, + ) + + /** + * Capabilities related to tools. + * + * @property listChanged Indicates if the server supports notifications when the tool list changes. + */ + @Serializable + data class Tools( + /** + * Whether this server supports issuing notifications for changes to the tool list. + */ + val listChanged: Boolean?, + ) +} + +/** + * After receiving an initialized request from the client, the server sends this response. + */ +@Serializable +data class InitializeResult( + /** + * The version of the Model Context Protocol that the server wants to use. This may not match the version that the client requested. If the client cannot support this version, it MUST disconnect. + */ + val protocolVersion: String = LATEST_PROTOCOL_VERSION, + val capabilities: ServerCapabilities = ServerCapabilities(), + val serverInfo: Implementation, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult + +/** + * This notification is sent from the client to the server after initialization has finished. + */ +@Serializable +class InitializedNotification : ClientNotification { + override val method: Method = Method.Defined.NotificationsInitialized +} + +/* Ping */ +/** + * A ping, issued by either the server or the client, to check that the other party is still alive. The receiver must promptly respond, or else may be disconnected. + */ +@Serializable +class PingRequest : ServerRequest, ClientRequest { + override val method: Method = Method.Defined.Ping +} + +/** + * Represents the base interface for progress tracking. + */ +@Serializable +sealed interface ProgressBase { + /** + * The progress thus far. This should increase every time progress is made, even if the total is unknown. + */ + val progress: Int + + /** + * Total number of items to a process (or total progress required), if known. + */ + val total: Double? +} + +/* org.jetbrains.kotlinx.mcp.Progress notifications */ +/** + * Represents a progress notification. + * + * @property progress The current progress value. + * @property total The total progress required, if known. + */ +@Serializable +open class Progress( + /** + * The progress thus far. This should increase every time progress is made, even if the total is unknown. + */ + override val progress: Int, + + /** + * Total number of items to a process (or total progress required), if known. + */ + override val total: Double?, +) : ProgressBase + +/** + * An out-of-band notification used to inform the receiver of a progress update for a long-running request. + */ +@Serializable +data class ProgressNotification( + override val progress: Int, + /** + * The progress token, + * which was given in the initial request, + * used to associate this notification with the request that is proceeding. + */ + val progressToken: ProgressToken, + @Suppress("PropertyName") val _meta: JsonObject = EmptyJsonObject, + override val total: Double?, +) : ClientNotification, ServerNotification, ProgressBase { + override val method: Method = Method.Defined.NotificationsProgress +} + +/* Pagination */ +/** + * Represents a request supporting pagination. + */ +@Serializable +sealed interface PaginatedRequest : Request, WithMeta { + /** + * The cursor indicating the pagination position. + */ + val cursor: Cursor? + override val _meta: JsonObject +} + +/** + * Represents a paginated result of a request. + */ +@Serializable +sealed interface PaginatedResult : RequestResult { + /** + * An opaque token representing the pagination position after the last returned result. + * If present, there may be more results available. + */ + val nextCursor: Cursor? +} + +/* Resources */ +/** + * The contents of a specific resource or sub-resource. + */ +@Serializable(with = ResourceContentsPolymorphicSerializer::class) +sealed interface ResourceContents { + /** + * The URI of this resource. + */ + val uri: String + + /** + * The MIME type of this resource, if known. + */ + val mimeType: String? +} + +/** + * Represents the text contents of a resource. + * + * @property text The text of the item. This must only be set if the item can actually be represented as text (not binary data). + */ +@Serializable +data class TextResourceContents( + val text: String, + override val uri: String, + override val mimeType: String?, +) : ResourceContents + +/** + * Represents the binary contents of a resource encoded as a base64 string. + * + * @property blob A base64-encoded string representing the binary data of the item. + */ +@Serializable +data class BlobResourceContents( + val blob: String, + override val uri: String, + override val mimeType: String?, +) : ResourceContents + +/** + * Represents resource contents with unknown or unspecified data. + */ +@Serializable +data class UnknownResourceContents( + override val uri: String, + override val mimeType: String?, +) : ResourceContents + +/** + * A known resource that the server is capable of reading. + */ +@Serializable +data class Resource( + /** + * The URI of this resource. + */ + val uri: String, + /** + * A human-readable name for this resource. + * + * Clients can use this to populate UI elements. + */ + val name: String, + /** + * A description of what this resource represents. + * + * Clients can use this to improve the LLM's understanding of available resources. + * It can be thought of as a "hint" to the model. + */ + val description: String?, + /** + * The MIME type of this resource, if known. + */ + val mimeType: String?, +) + +/** + * A template description for resources available on the server. + */ +@Serializable +data class ResourceTemplate( + /** + * A URI template (according to RFC 6570) that can be used to construct resource URIs. + */ + val uriTemplate: String, + /** + * A human-readable name for the type of resource this template refers to. + * + * Clients can use this to populate UI elements. + */ + val name: String, + /** + * A description of what this template is for. + * + * Clients can use this to improve the LLM's understanding of available resources. + * It can be thought of as a "hint" to the model. + */ + val description: String?, + /** + * The MIME type for all resources that match this template. This should only be included if all resources matching this template have the same type. + */ + val mimeType: String?, +) + +/** + * Sent from the client to request a list of resources the server has. + */ +@Serializable +data class ListResourcesRequest( + override val cursor: Cursor? = null, + override val _meta: JsonObject = EmptyJsonObject +) : ClientRequest, PaginatedRequest { + override val method: Method = Method.Defined.ResourcesList +} + +/** + * The server's response to a resources/list request from the client. + */ +@Serializable +class ListResourcesResult( + val resources: List, + override val nextCursor: Cursor? = null, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult, PaginatedResult + +/** + * Sent from the client to request a list of resource templates the server has. + */ +@Serializable +data class ListResourceTemplatesRequest( + override val cursor: Cursor?, + override val _meta: JsonObject = EmptyJsonObject +) : ClientRequest, PaginatedRequest { + override val method: Method = Method.Defined.ResourcesTemplatesList +} + +/** + * The server's response to a resources/templates/list request from the client. + */ +@Serializable +class ListResourceTemplatesResult( + val resourceTemplates: List, + override val nextCursor: Cursor? = null, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult, PaginatedResult + +/** + * Sent from the client to the server to read a specific resource URI. + */ +@Serializable +data class ReadResourceRequest( + val uri: String, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.ResourcesRead +} + +/** + * The server's response to a resources/read request from the client. + */ +@Serializable +class ReadResourceResult( + val contents: List, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult + +/** + * An optional notification from the server to the client, + * informing it that the list of resources it can read from has changed. + * Servers may issue this without any previous subscription from the client. + */ +@Serializable +class ResourceListChangedNotification : ServerNotification { + override val method: Method = Method.Defined.NotificationsResourcesListChanged +} + +/** + * Sent from the client to request resources/updated notifications from the server whenever a particular resource changes. + */ +@Serializable +data class SubscribeRequest( + /** + * The URI of the resource to subscribe to. The URI can use any protocol; it is up to the server how to interpret it. + */ + val uri: String, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.ResourcesSubscribe +} + +/** + * Sent from the client to request cancellation of resources/updated notifications from the server. This should follow a previous resources/subscribe request. + */ +@Serializable +data class UnsubscribeRequest( + /** + * The URI of the resource to unsubscribe from. + */ + val uri: String, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.ResourcesUnsubscribe +} + +/** + * A notification from the server to the client, informing it that a resource has changed and may need to be read again. This should only be sent if the client previously sent a resources/subscribe request. + */ +@Serializable +data class ResourceUpdatedNotification( + /** + * The URI of the resource that has been updated. This might be a sub-resource of the one that the client actually subscribed to. + */ + val uri: String, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerNotification, WithMeta { + override val method: Method = Method.Defined.NotificationsResourcesUpdated +} + +/* Prompts */ +/** + * Describes an argument that a prompt can accept. + */ +@Serializable +data class PromptArgument( + /** + * The name of the argument. + */ + val name: String, + /** + * A human-readable description of the argument. + */ + val description: String?, + /** + * Whether this argument must be provided. + */ + val required: Boolean?, +) + +/** + * A prompt or prompt template that the server offers. + */ +@Serializable +class Prompt( + /** + * The name of the prompt or prompt template. + */ + val name: String, + /** + * An optional description of what this prompt provides + */ + val description: String?, + /** + * A list of arguments to use for templating the prompt. + */ + val arguments: List?, +) + +/** + * Sent from the client to request a list of prompts and prompt templates the server has. + */ +@Serializable +data class ListPromptsRequest( + override val cursor: Cursor? = null, + override val _meta: JsonObject = EmptyJsonObject +) : ClientRequest, PaginatedRequest { + override val method: Method = Method.Defined.PromptsList +} + +/** + * The server's response to a prompts/list request from the client. + */ +@Serializable +class ListPromptsResult( + val prompts: List, + override val nextCursor: Cursor? = null, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult, PaginatedResult + +/** + * Used by the client to get a prompt provided by the server. + */ +@Serializable +data class GetPromptRequest( + /** + * The name of the prompt or prompt template. + */ + val name: String, + + /** + * Arguments to use for templating the prompt. + */ + val arguments: Map?, + + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.PromptsGet +} + +/** + * Represents the content of a prompt message. + */ +@Serializable(with = PromptMessageContentPolymorphicSerializer::class) +sealed interface PromptMessageContent { + val type: String +} + +/** + * Represents prompt message content that is either text or an image. + */ +@Serializable(with = PromptMessageContentTextOrImagePolymorphicSerializer::class) +sealed interface PromptMessageContentTextOrImage : PromptMessageContent + +/** + * Text provided to or from an LLM. + */ +@Serializable +data class TextContent( + /** + * The text content of the message. + */ + val text: String? = null, +) : PromptMessageContentTextOrImage { + override val type: String = TYPE + + companion object { + const val TYPE = "text" + } +} + +/** + * An image provided to or from an LLM. + */ +@Serializable +data class ImageContent( + /** + * The base64-encoded image data. + */ + val data: String, + + /** + * The MIME type of the image. Different providers may support different image types. + */ + val mimeType: String, +) : PromptMessageContentTextOrImage { + override val type: String = TYPE + + companion object { + const val TYPE = "image" + } +} + +/** + * An image provided to or from an LLM. + */ +@Serializable +data class UnknownContent( + override val type: String, +) : PromptMessageContentTextOrImage + +/** + * The contents of a resource, embedded into a prompt or tool call result. + */ +@Serializable +data class EmbeddedResource( + val resource: ResourceContents, +) : PromptMessageContent { + override val type: String = TYPE + + companion object { + const val TYPE = "resource" + } +} + +/** + * Enum representing the role of a participant. + */ +@Suppress("EnumEntryName") +@Serializable +enum class Role { + user, assistant, +} + +/** + * Describes a message returned as part of a prompt. + */ +@Serializable +data class PromptMessage( + val role: Role, + val content: PromptMessageContent, +) + +/** + * The server's response to a prompts/get request from the client. + */ +@Serializable +class GetPromptResult( + /** + * An optional description for the prompt. + */ + val description: String?, + val messages: List, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult + +/** + * An optional notification from the server to the client, informing it that the list of prompts it offers has changed. + * Servers may issue this without any previous subscription from the client. + */ +@Serializable +class PromptListChangedNotification : ServerNotification { + override val method: Method = Method.Defined.NotificationsPromptsListChanged +} + +/* Tools */ +/** + * Definition for a tool the client can call. + */ +@Serializable +data class Tool( + /** + * The name of the tool. + */ + val name: String, + /** + * A human-readable description of the tool. + */ + val description: String?, + /** + * A JSON object defining the expected parameters for the tool. + */ + val inputSchema: Input, +) { + @Serializable + data class Input( + val properties: JsonObject = EmptyJsonObject, + ) { + val type: String = "object" + } +} + +/** + * Sent from the client to request a list of tools the server has. + */ +@Serializable +data class ListToolsRequest( + override val cursor: Cursor? = null, + override val _meta: JsonObject = EmptyJsonObject +) : ClientRequest, PaginatedRequest { + override val method: Method = Method.Defined.ToolsList +} + +/** + * The server's response to a tools/list request from the client. + */ +@Serializable +class ListToolsResult( + val tools: List, + override val nextCursor: Cursor?, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult, PaginatedResult + +/** + * The server's response to a tool call. + */ +@Serializable +sealed interface CallToolResultBase : ServerResult { + val content: List + val isError: Boolean? get() = false +} + +/** + * The server's response to a tool call. + */ +@Serializable +class CallToolResult( + override val content: List, + override val isError: Boolean? = false, + override val _meta: JsonObject = EmptyJsonObject, +) : CallToolResultBase + +/** + * [CallToolResult] extended with backwards compatibility to protocol version 2024-10-07. + */ +@Serializable +class CompatibilityCallToolResult( + override val content: List, + override val isError: Boolean? = false, + override val _meta: JsonObject = EmptyJsonObject, + val toolResult: JsonObject = EmptyJsonObject, +) : CallToolResultBase + +/** + * Used by the client to invoke a tool provided by the server. + */ +@Serializable +data class CallToolRequest( + val name: String, + val arguments: JsonObject = EmptyJsonObject, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.ToolsCall +} + +/** + * An optional notification from the server to the client, informing it that the list of tools it offers has changed. + * Servers may issue this without any previous subscription from the client. + */ +@Serializable +class ToolListChangedNotification : ServerNotification { + override val method: Method = Method.Defined.NotificationsToolsListChanged +} + +/* Logging */ +/** + * The severity of a log message. + */ +@Suppress("EnumEntryName") +@Serializable +enum class LoggingLevel { + debug, + info, + notice, + warning, + error, + critical, + alert, + emergency, + ; +} + +/** + * org.jetbrains.kotlinx.mcp.Notification of a log message passed from server to client. + * If no logging level request has been sent from the client, + * the server MAY decide which messages to send automatically. + */ +@Serializable +data class LoggingMessageNotification( + /** + * The severity of this log message. + */ + val level: LoggingLevel, + + /** + * An optional name of the logger issuing this message. + */ + val logger: String? = null, + + /** + * The data to be logged, such as a string message or an object. Any JSON serializable type is allowed here. + */ + val data: JsonObject = EmptyJsonObject, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerNotification, WithMeta { + /** + * A request from the client to the server to enable or adjust logging. + */ + @Serializable + data class SetLevelRequest( + /** + * The level of logging that the client wants to receive from the server. The server should send all logs at this level and higher (i.e., more severe) to the client as notifications/logging/message. + */ + val level: LoggingLevel, + override val _meta: JsonObject = EmptyJsonObject, + ) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.LoggingSetLevel + } + + override val method: Method = Method.Defined.NotificationsMessage +} + +/* Sampling */ +/** + * Hints to use for model selection. + */ +@Serializable +data class ModelHint( + /** + * A hint for a model name. + */ + val name: String?, +) + +/** + * The server's preferences for model selection, requested by the client during sampling. + */ +@Suppress("CanBeParameter") +@Serializable +class ModelPreferences( + /** + * Optional hints to use for model selection. + */ + val hints: List?, + /** + * How much to prioritize cost when selecting a model. + */ + val costPriority: Double?, + /** + * How much to prioritize sampling speed (latency) when selecting a model. + */ + val speedPriority: Double?, + /** + * How much to prioritize intelligence and capabilities when selecting a model. + */ + val intelligencePriority: Double?, +) { + init { + require(costPriority == null || costPriority in 0.0..1.0) { + "costPriority must be in 0.0 <= x <= 1.0 value range" + } + + require(speedPriority == null || speedPriority in 0.0..1.0) { + "costPriority must be in 0.0 <= x <= 1.0 value range" + } + + require(intelligencePriority == null || intelligencePriority in 0.0..1.0) { + "intelligencePriority must be in 0.0 <= x <= 1.0 value range" + } + } +} + +/** + * Describes a message issued to or received from an LLM API. + */ +@Serializable +data class SamplingMessage( + val role: Role, + val content: PromptMessageContentTextOrImage, +) + +/** + * A request from the server to sample an LLM via the client. + * The client has full discretion over which model to select. + * The client should also inform the user before beginning sampling to allow them to inspect the request + * (human in the loop) and decide whether to approve it. + */ +@Serializable +data class CreateMessageRequest( + val messages: List, + /** + * An optional system prompt the server wants to use it for sampling. The client MAY modify or omit this prompt. + */ + val systemPrompt: String?, + /** + * A request to include context from one or more MCP servers (including the caller), to be attached to the prompt. The client MAY ignore this request. + */ + val includeContext: IncludeContext?, + val temperature: Double?, + /** + * The maximum number of tokens to sample, as requested by the server. The client MAY choose to sample fewer tokens than requested. + */ + val maxTokens: Int, + val stopSequences: List?, + /** + * Optional metadata to pass through to the LLM provider. The format of this metadata is provider-specific. + */ + val metadata: JsonObject = EmptyJsonObject, + /** + * The server's preferences for which model to select. + */ + val modelPreferences: ModelPreferences?, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerRequest, WithMeta { + override val method: Method = Method.Defined.SamplingCreateMessage + + @Serializable + enum class IncludeContext { none, thisServer, allServers } +} + +@Serializable(with = StopReasonSerializer::class) +sealed interface StopReason { + val value: String + + @Serializable + data object EndTurn : StopReason { + override val value: String = "endTurn" + } + + @Serializable + data object StopSequence : StopReason { + override val value: String = "stopSequence" + } + + @Serializable + data object MaxTokens : StopReason { + override val value: String = "maxTokens" + } + + @Serializable + @JvmInline + value class Other(override val value: String) : StopReason +} + +/** + * The client's response to a sampling/create_message request from the server. + * The client should inform the user before returning the sampled message to allow them to inspect the response + * (human in the loop) and decide whether to allow the server to see it. + */ +@Serializable +data class CreateMessageResult( + /** + * The name of the model that generated the message. + */ + val model: String, + /** + * The reason why sampling stopped. + */ + val stopReason: StopReason? = null, + val role: Role, + val content: PromptMessageContentTextOrImage, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientResult + +/* Autocomplete */ +@Serializable(with = ReferencePolymorphicSerializer::class) +sealed interface Reference { + val type: String +} + +/** + * A reference to a resource or resource template definition. + */ +@Serializable +data class ResourceReference( + /** + * The URI or URI template of the resource. + */ + val uri: String, +) : Reference { + override val type: String = TYPE + + companion object { + const val TYPE = "ref/resource" + } +} + +/** + * Identifies a prompt. + */ +@Serializable +data class PromptReference( + /** + * The name of the prompt or prompt template + */ + val name: String, +) : Reference { + override val type: String = TYPE + + companion object { + const val TYPE = "ref/prompt" + } +} + +/** + * Identifies a prompt. + */ +@Serializable +data class UnknownReference( + override val type: String, +) : Reference + +/** + * A request from the client to the server to ask for completion options. + */ +@Serializable +data class CompleteRequest( + val ref: Reference, + /** + * The argument's information + */ + val argument: Argument, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientRequest, WithMeta { + override val method: Method = Method.Defined.CompletionComplete + + @Serializable + data class Argument( + /** + * The name of the argument + */ + val name: String, + /** + * The value of the argument to use for completion matching. + */ + val value: String, + ) +} + +/** + * The server's response to a completion/complete request + */ +@Serializable +data class CompleteResult( + val completion: Completion, + override val _meta: JsonObject = EmptyJsonObject, +) : ServerResult { + @Suppress("CanBeParameter") + @Serializable + class Completion( + /** + * A list of completion values. Must not exceed 100 items. + */ + val values: List, + /** + * The total number of completion options available. This can exceed the number of values actually sent in the response. + */ + val total: Int?, + /** + * Indicates whether there are additional completion options beyond those provided in the current response, even if the exact total is unknown. + */ + val hasMore: Boolean?, + ) { + init { + require(values.size <= 100) { + "'values' field must not exceed 100 items" + } + } + } +} + +/* Roots */ +/** + * Represents a root directory or file that the server can operate on. + */ +@Serializable +data class Root( + /** + * The URI identifying the root. This *must* start with file:// for now. + */ + val uri: String, + + /** + * An optional name for the root. + */ + val name: String?, +) { + init { + require(uri.startsWith("file://")) { + "'uri' param must start with 'file://'" + } + } +} + +/** + * Sent from the server to request a list of root URIs from the client. + */ +@Serializable +class ListRootsRequest(override val _meta: JsonObject = EmptyJsonObject) : ServerRequest, WithMeta { + override val method: Method = Method.Defined.RootsList +} + +/** + * The client's response to a roots/list request from the server. + */ +@Serializable +class ListRootsResult( + val roots: List, + override val _meta: JsonObject = EmptyJsonObject, +) : ClientResult + +/** + * A notification from the client to the server, informing it that the list of roots has changed. + */ +@Serializable +class RootsListChangedNotification : ClientNotification { + override val method: Method = Method.Defined.NotificationsRootsListChanged +} + +/** + * Represents an error specific to the MCP protocol. + * + * @property code The error code. + * @property message The error message. + * @property data Additional error data as a JSON object. + */ +class McpError(val code: Int, message: String, val data: JsonObject = EmptyJsonObject) : Exception() { + override val message: String = "MCP error ${code}: $message" +} diff --git a/src/main/kotlin/org/jetbrains/kotlinx/mcp/types.util.kt b/src/main/kotlin/org/jetbrains/kotlinx/mcp/types.util.kt new file mode 100644 index 0000000..1b6a539 --- /dev/null +++ b/src/main/kotlin/org/jetbrains/kotlinx/mcp/types.util.kt @@ -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 { + 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 { + 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 { + 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::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return when (element.jsonObject.getValue("type").jsonPrimitive.content) { + ResourceReference.TYPE -> ResourceReference.serializer() + PromptReference.TYPE -> PromptReference.serializer() + else -> UnknownReference.serializer() + } + } +} + +internal object PromptMessageContentPolymorphicSerializer : + JsonContentPolymorphicSerializer(PromptMessageContent::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + 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::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return when (element.jsonObject.getValue("type").jsonPrimitive.content) { + ImageContent.TYPE -> ImageContent.serializer() + TextContent.TYPE -> TextContent.serializer() + else -> UnknownContent.serializer() + } + } +} + +internal object ResourceContentsPolymorphicSerializer : + JsonContentPolymorphicSerializer(ResourceContents::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + 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 { + selectClientRequestDeserializer(method)?.let { return it } + selectServerRequestDeserializer(method)?.let { return it } + return CustomRequest.serializer() +} + +internal fun selectClientRequestDeserializer(method: String): DeserializationStrategy? { + 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::class) { +// override fun selectDeserializer(element: JsonElement): DeserializationStrategy { +// 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? { + 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::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return selectClientNotificationDeserializer(element) + ?: UnknownMethodRequestOrNotification.serializer() + } +} + +internal fun selectServerRequestDeserializer(method: String): DeserializationStrategy? { + 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::class) { +// override fun selectDeserializer(element: JsonElement): DeserializationStrategy { +// return org.jetbrains.kotlinx.mcp.selectServerRequestDeserializer(element) +// ?: org.jetbrains.kotlinx.mcp.UnknownMethodRequestOrNotification.serializer() +// } +//} + +internal fun selectServerNotificationDeserializer(element: JsonElement): DeserializationStrategy? { + 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::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return selectServerNotificationDeserializer(element) + ?: UnknownMethodRequestOrNotification.serializer() + } +} + +internal object NotificationPolymorphicSerializer : + JsonContentPolymorphicSerializer(Notification::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return selectClientNotificationDeserializer(element) + ?: selectServerNotificationDeserializer(element) + ?: UnknownMethodRequestOrNotification.serializer() + } +} + +internal object RequestPolymorphicSerializer : + JsonContentPolymorphicSerializer(Request::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + 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, + * total: Int?, + * hasMore: Boolean?, + * } + * _meta: JsonObject, + * } + * + * org.jetbrains.kotlinx.mcp.GetPromptResult { + * description: String?, + * messages: Array, + * _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, + * nextCursor: org.jetbrains.kotlinx.mcp.Cursor?, + * _meta: JsonObject, + * } + * + * org.jetbrains.kotlinx.mcp.ListResourceTemplatesResult { + * resourceTemplates: Array, + * nextCursor: org.jetbrains.kotlinx.mcp.Cursor?, + * _meta: JsonObject, + * } + * + * org.jetbrains.kotlinx.mcp.ListResourcesResult { + * resources: Array, + * nextCursor: org.jetbrains.kotlinx.mcp.Cursor?, + * _meta: JsonObject, + * } + * + * org.jetbrains.kotlinx.mcp.ListToolsResult { + * tools: Array, + * nextCursor: org.jetbrains.kotlinx.mcp.Cursor?, + * _meta: JsonObject, + * } + * + * org.jetbrains.kotlinx.mcp.ReadResourceResult { + * contents: Array, + * _meta: JsonObject, + * } + * ``` + */ +private fun selectServerResultDeserializer(element: JsonElement): DeserializationStrategy? { + 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, + * _meta: JsonObject, + * } + *``` + */ +private fun selectClientResultDeserializer(element: JsonElement): DeserializationStrategy? { + val jsonObject = element.jsonObject + return when { + jsonObject.contains("model") -> CreateMessageResult.serializer() + jsonObject.contains("roots") -> ListRootsResult.serializer() + else -> null + } +} + +internal object ServerResultPolymorphicSerializer : + JsonContentPolymorphicSerializer(ServerResult::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return selectServerResultDeserializer(element) + ?: EmptyRequestResult.serializer() + } +} + +internal object ClientResultPolymorphicSerializer : + JsonContentPolymorphicSerializer(ClientResult::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return selectClientResultDeserializer(element) + ?: EmptyRequestResult.serializer() + } +} + +internal object RequestResultPolymorphicSerializer : + JsonContentPolymorphicSerializer(RequestResult::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + return selectClientResultDeserializer(element) + ?: selectServerResultDeserializer(element) + ?: EmptyRequestResult.serializer() + } +} + +internal object JSONRPCMessagePolymorphicSerializer : + JsonContentPolymorphicSerializer(JSONRPCMessage::class) { + override fun selectDeserializer(element: JsonElement): DeserializationStrategy { + 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()) diff --git a/src/test/kotlin/InMemoryTransport.kt b/src/test/kotlin/InMemoryTransport.kt new file mode 100644 index 0000000..01d832b --- /dev/null +++ b/src/test/kotlin/InMemoryTransport.kt @@ -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 = 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 { + 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) + } + } +} \ No newline at end of file diff --git a/src/test/kotlin/client/BaseTransportTest.kt b/src/test/kotlin/client/BaseTransportTest.kt new file mode 100644 index 0000000..db3fbec --- /dev/null +++ b/src/test/kotlin/client/BaseTransportTest.kt @@ -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( + PingRequest().toJSON(), + InitializedNotification().toJSON() + ) + + val readMessages = mutableListOf() + val finished = CompletableDeferred() + + 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() + } +} diff --git a/src/test/kotlin/client/ClientIntegrationTest.kt b/src/test/kotlin/client/ClientIntegrationTest.kt new file mode 100644 index 0000000..00a8684 --- /dev/null +++ b/src/test/kotlin/client/ClientIntegrationTest.kt @@ -0,0 +1,39 @@ +package client + +import org.jetbrains.kotlinx.mcp.Implementation +import org.jetbrains.kotlinx.mcp.ListToolsResult +import kotlinx.coroutines.test.runTest +import org.jetbrains.kotlinx.mcp.client.Client +import org.jetbrains.kotlinx.mcp.client.StdioClientTransport +import org.junit.jupiter.api.Disabled +import org.junit.jupiter.api.Test +import java.net.Socket + +class ClientIntegrationTest { + + fun createTransport(): StdioClientTransport { + val socket = Socket("localhost", 3000) + + return StdioClientTransport(socket.inputStream, socket.outputStream) + } + + @Disabled("This test requires a running server") + @Test + fun testRequestTools() = runTest { + val client = Client( + Implementation("test", "1.0"), + ) + + val transport = createTransport() + try { + client.connect(transport) + + val response: ListToolsResult? = client.listTools() + println(response?.tools) + + } finally { + transport.close() + } + } + +} diff --git a/src/test/kotlin/client/ClientTest.kt b/src/test/kotlin/client/ClientTest.kt new file mode 100644 index 0000000..3ff771f --- /dev/null +++ b/src/test/kotlin/client/ClientTest.kt @@ -0,0 +1,510 @@ +package client + +import org.jetbrains.kotlinx.mcp.ClientCapabilities +import org.jetbrains.kotlinx.mcp.CreateMessageRequest +import org.jetbrains.kotlinx.mcp.CreateMessageResult +import org.jetbrains.kotlinx.mcp.EmptyJsonObject +import org.jetbrains.kotlinx.mcp.Implementation +import InMemoryTransport +import org.jetbrains.kotlinx.mcp.InitializeRequest +import org.jetbrains.kotlinx.mcp.InitializeResult +import org.jetbrains.kotlinx.mcp.JSONRPCMessage +import org.jetbrains.kotlinx.mcp.JSONRPCRequest +import org.jetbrains.kotlinx.mcp.JSONRPCResponse +import org.jetbrains.kotlinx.mcp.LATEST_PROTOCOL_VERSION +import org.jetbrains.kotlinx.mcp.ListResourcesRequest +import org.jetbrains.kotlinx.mcp.ListResourcesResult +import org.jetbrains.kotlinx.mcp.ListRootsRequest +import org.jetbrains.kotlinx.mcp.ListToolsRequest +import org.jetbrains.kotlinx.mcp.ListToolsResult +import org.jetbrains.kotlinx.mcp.LoggingLevel +import org.jetbrains.kotlinx.mcp.LoggingMessageNotification +import org.jetbrains.kotlinx.mcp.Method +import org.jetbrains.kotlinx.mcp.Role +import org.jetbrains.kotlinx.mcp.SUPPORTED_PROTOCOL_VERSIONS +import org.jetbrains.kotlinx.mcp.ServerCapabilities +import org.jetbrains.kotlinx.mcp.TextContent +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.TimeoutCancellationException +import kotlinx.coroutines.cancel +import kotlinx.coroutines.joinAll +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlinx.serialization.json.buildJsonObject +import kotlinx.serialization.json.put +import org.jetbrains.kotlinx.mcp.client.Client +import org.jetbrains.kotlinx.mcp.client.ClientOptions +import org.junit.jupiter.api.Test +import org.jetbrains.kotlinx.mcp.server.Server +import org.jetbrains.kotlinx.mcp.server.ServerOptions +import org.jetbrains.kotlinx.mcp.shared.Transport +import kotlin.coroutines.cancellation.CancellationException +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import kotlin.test.fail + +class ClientTest { + @Test + fun `should initialize with matching protocol version`() = runTest { + var initialied = false + var clientTransport = object : Transport { + override suspend fun start() {} + + override suspend fun send(message: JSONRPCMessage) { + if (message !is JSONRPCRequest) return + initialied = true + val result = InitializeResult( + protocolVersion = LATEST_PROTOCOL_VERSION, + capabilities = ServerCapabilities(), + serverInfo = Implementation( + name = "test", + version = "1.0" + ) + ) + + val response = JSONRPCResponse( + id = message.id, + result = result + ) + + onMessage?.invoke(response) + } + + override suspend fun close() { + } + + override var onClose: (() -> Unit)? = null + override var onError: ((Throwable) -> Unit)? = null + override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null + } + + val client = Client( + clientInfo = Implementation( + name = "test client", + version = "1.0" + ), + options = ClientOptions( + capabilities = ClientCapabilities( + sampling = EmptyJsonObject + ) + ) + ) + + client.connect(clientTransport) + assertTrue(initialied) + } + + @Test + fun `should initialize with supported older protocol version`() = runTest { + val OLD_VERSION = SUPPORTED_PROTOCOL_VERSIONS[1] + var clientTransport = object : Transport { + override suspend fun start() {} + + override suspend fun send(message: JSONRPCMessage) { + if (message !is JSONRPCRequest) return + check(message.method == Method.Defined.Initialize.value) + + val result = InitializeResult( + protocolVersion = OLD_VERSION, + capabilities = ServerCapabilities(), + serverInfo = Implementation( + name = "test", + version = "1.0" + ) + ) + + val response = JSONRPCResponse( + id = message.id, + result = result + ) + onMessage?.invoke(response) + } + + override suspend fun close() { + } + + override var onClose: (() -> Unit)? = null + override var onError: ((Throwable) -> Unit)? = null + override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null + } + + val client = Client( + clientInfo = Implementation( + name = "test client", + version = "1.0" + ), + options = ClientOptions( + capabilities = ClientCapabilities( + sampling = EmptyJsonObject + ) + ) + ) + + client.connect(clientTransport) + assertEquals( + Implementation("test", "1.0"), + client.getServerVersion() + ) + } + + @Test + fun `should reject unsupported protocol version`() = runTest { + var closed = false + val clientTransport = object : Transport { + override suspend fun start() {} + + override suspend fun send(message: JSONRPCMessage) { + if (message !is JSONRPCRequest) return + check(message.method == Method.Defined.Initialize.value) + + val result = InitializeResult( + protocolVersion = "invalid-version", + capabilities = ServerCapabilities(), + serverInfo = Implementation( + name = "test", + version = "1.0" + ) + ) + + val response = JSONRPCResponse( + id = message.id, + result = result + ) + + onMessage?.invoke(response) + } + + override suspend fun close() { + closed = true + } + + override var onClose: (() -> Unit)? = null + override var onError: ((Throwable) -> Unit)? = null + override var onMessage: (suspend (JSONRPCMessage) -> Unit)? = null + } + + val client = Client( + clientInfo = Implementation( + name = "test client", + version = "1.0" + ), + options = ClientOptions() + ) + + assertFailsWith("Server's protocol version is not supported: invalid-version") { + client.connect(clientTransport) + } + + assertTrue(closed) + } + + @Test + fun `should respect server capabilities`() = runTest { + val serverOptions = ServerOptions( + capabilities = ServerCapabilities( + resources = ServerCapabilities.Resources(null, null), + tools = ServerCapabilities.Tools(null) + ) + ) + val server = Server( + Implementation(name = "test server", version = "1.0"), + serverOptions + ) + + server.setRequestHandler(Method.Defined.Initialize) { request, _ -> + InitializeResult( + protocolVersion = LATEST_PROTOCOL_VERSION, + capabilities = ServerCapabilities( + resources = ServerCapabilities.Resources(null, null), + tools = ServerCapabilities.Tools(null) + ), + serverInfo = Implementation(name = "test", version = "1.0") + ) + } + + server.setRequestHandler(Method.Defined.ResourcesList) { request, _ -> + ListResourcesResult(resources = emptyList(), nextCursor = null) + } + + server.setRequestHandler(Method.Defined.ToolsList) { request, _ -> + ListToolsResult(tools = emptyList(), nextCursor = null) + } + + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + + val client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions( + capabilities = ClientCapabilities(sampling = EmptyJsonObject), + ) + ) + + listOf( + launch { + client.connect(clientTransport) + }, + launch { + server.connect(serverTransport) + } + ).joinAll() + + // Server supports resources and tools, but not prompts + val caps = client.getServerCapabilities() + assertEquals(ServerCapabilities.Resources(null, null), caps?.resources) + assertEquals(ServerCapabilities.Tools(null), caps?.tools) + assertTrue(caps?.prompts == null) // or check that prompts are absent + + // These should not throw + client.listResources() + client.listTools() + + // This should fail because prompts are not supported + val ex = assertFailsWith { + client.listPrompts() + } + assertTrue(ex.message?.contains("Server does not support prompts") == true) + } + + @Test + fun `should respect client notification capabilities`() = runTest { + val server = Server( + Implementation(name = "test server", version = "1.0"), + ServerOptions(capabilities = ServerCapabilities()) + ) + + val client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions( + capabilities = ClientCapabilities( + roots = ClientCapabilities.Roots(listChanged = true) + ) + ) + ) + + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + + listOf( + launch { + client.connect(clientTransport) + println("Client connected") + }, + launch { + server.connect(serverTransport) + println("Server connected") + } + ).joinAll() + + // This should not throw because the client supports roots.listChanged + client.sendRootsListChanged() + + // Create a new client without the roots.listChanged capability + val clientWithoutCapability = Client( + clientInfo = Implementation(name = "test client without capability", version = "1.0"), + options = ClientOptions( + capabilities = ClientCapabilities(), + // enforceStrictCapabilities = true // TODO() + ) + ) + + clientWithoutCapability.connect(clientTransport) + // Using the same transport pair might not be realistic - in a real scenario you'd create another pair. + // Adjust if necessary. + + // This should fail + val ex = assertFailsWith { + clientWithoutCapability.sendRootsListChanged() + } + assertTrue(ex.message?.startsWith("Client does not support") == true) + } + + @Test + fun `should respect server notification capabilities`() = runTest { + val server = Server( + Implementation(name = "test server", version = "1.0"), + ServerOptions( + capabilities = ServerCapabilities( + logging = EmptyJsonObject, + resources = ServerCapabilities.Resources(listChanged = true, subscribe = null) + ) + ) + ) + + val client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions( + capabilities = ClientCapabilities() + ) + ) + + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + + listOf( + launch { + client.connect(clientTransport) + println("Client connected") + }, + launch { + server.connect(serverTransport) + println("Server connected") + } + ).joinAll() + + // These should not throw + val jsonObject = buildJsonObject { + put("name", "John") + put("age", 30) + put("isStudent", false) + } + server.sendLoggingMessage( + LoggingMessageNotification( + level = LoggingLevel.info, + data = jsonObject + ) + ) + server.sendResourceListChanged() + + // This should fail because the server doesn't have the tools capability + val ex = assertFailsWith { + server.sendToolListChanged() + } + assertTrue(ex.message?.contains("Server does not support notifying of tool list changes") == true) + } + + @Test + fun `should handle client cancelling a request`() = runTest { + val server = Server( + Implementation(name = "test server", version = "1.0"), + ServerOptions( + capabilities = ServerCapabilities(resources = ServerCapabilities.Resources(listChanged = null, subscribe = null)) + ) + ) + + val def = CompletableDeferred() + val defTimeOut = CompletableDeferred() + server.setRequestHandler(Method.Defined.ResourcesList) { _, extra -> + // Simulate delay + def.complete(Unit) + try { + kotlinx.coroutines.delay(1000) + } catch (e: CancellationException) { + defTimeOut.complete(Unit) + throw e + } + fail("Shouldn't have been called") + ListResourcesResult(resources = emptyList()) + } + + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + + val client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions(capabilities = ClientCapabilities()) + ) + + listOf( + launch { + client.connect(clientTransport) + println("Client connected") + }, + launch { + server.connect(serverTransport) + println("Server connected") + } + ).joinAll() + + val defCancel = CompletableDeferred() + val job = launch { + try { + client.listResources() + } catch (e: CancellationException) { + defCancel.complete(Unit) + assertEquals("Cancelled by test", e.message) + } + } + def.await() + runCatching { job.cancel("Cancelled by test") } + defCancel.await() + defTimeOut.await() + } + + @Test + fun `should handle request timeout`() = runTest { + val server = Server( + Implementation(name = "test server", version = "1.0"), + ServerOptions( + capabilities = ServerCapabilities(resources = ServerCapabilities.Resources(listChanged = null, subscribe = null)) + ) + ) + + server.setRequestHandler(Method.Defined.ResourcesList) { _, extra -> + // Simulate a delayed response + // Wait ~100ms unless cancelled + try { + kotlinx.coroutines.withTimeout(100L) { + // Just delay here, if timeout is 0 on client side this won't return in time + kotlinx.coroutines.delay(100) + } + } catch (e: Exception) { + // If aborted, just rethrow or return early + } + ListResourcesResult(resources = emptyList()) + } + + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + val client = Client( + clientInfo = Implementation(name = "test client", version = "1.0"), + options = ClientOptions(capabilities = ClientCapabilities()) + ) + + listOf( + launch { + client.connect(clientTransport) + println("Client connected") + }, + launch { + server.connect(serverTransport) + println("Server connected") + } + ).joinAll() + + // Request with 1 msec timeout should fail immediately + val ex = assertFailsWith { + kotlinx.coroutines.withTimeout(1) { + client.listResources() + } + } + assertTrue(ex is TimeoutCancellationException) + } + + @Test + fun `should only allow setRequestHandler for declared capabilities`() = runTest { + val client = Client( + clientInfo = Implementation( + name = "test client", + version = "1.0" + ), + options = ClientOptions( + capabilities = ClientCapabilities( + sampling = EmptyJsonObject + ) + ) + ) + + client.setRequestHandler(Method.Defined.SamplingCreateMessage) { request, _ -> + CreateMessageResult( + model = "test-model", + role = Role.assistant, + content = TextContent( + text = "Test response" + ) + ) + } + + assertFailsWith("Client does not support roots capability (required for RootsList)") { + client.setRequestHandler(Method.Defined.RootsList) { _, _ -> null } + } + } + + +} diff --git a/src/test/kotlin/client/InMemoryTransportTest.kt b/src/test/kotlin/client/InMemoryTransportTest.kt new file mode 100644 index 0000000..b21450c --- /dev/null +++ b/src/test/kotlin/client/InMemoryTransportTest.kt @@ -0,0 +1,124 @@ +package client + +import InMemoryTransport +import org.jetbrains.kotlinx.mcp.InitializedNotification +import org.jetbrains.kotlinx.mcp.JSONRPCMessage +import kotlinx.coroutines.runBlocking +import org.jetbrains.kotlinx.mcp.toJSON +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +class InMemoryTransportTest { + private lateinit var clientTransport: InMemoryTransport + private lateinit var serverTransport: InMemoryTransport + + @BeforeEach + fun setUp() { + val (client, server) = InMemoryTransport.createLinkedPair() + clientTransport = client + serverTransport = server + } + + @Test + fun `should create linked pair`() { + assertNotNull(clientTransport) + assertNotNull(serverTransport) + } + + @Test + fun `should start without error`() { + runBlocking { + clientTransport.start() + serverTransport.start() + // If no exception is thrown, the test passes + } + } + + @Test + fun `should send message from client to server`() { + runBlocking { + val message = InitializedNotification() + + var receivedMessage: JSONRPCMessage? = null + serverTransport.onMessage = { msg -> + receivedMessage = msg + } + + val rpcNotification = message.toJSON() + clientTransport.send(rpcNotification) + assertEquals(rpcNotification, receivedMessage) + } + } + + @Test + fun `should send message from server to client`() { + runBlocking { + val message = InitializedNotification() + .toJSON() + + var receivedMessage: JSONRPCMessage? = null + clientTransport.onMessage = { msg -> + receivedMessage = msg + } + + serverTransport.send(message) + assertEquals(message, receivedMessage) + } + } + + @Test + fun `should handle close`() { + runBlocking { + var clientClosed = false + var serverClosed = false + + clientTransport.onClose = { + clientClosed = true + } + + serverTransport.onClose = { + serverClosed = true + } + + clientTransport.close() + assertTrue(clientClosed) + assertTrue(serverClosed) + } + } + + @Test + fun `should throw error when sending after close`() { + runBlocking { + clientTransport.close() + + assertThrows { + runBlocking { + clientTransport.send( + InitializedNotification().toJSON() + ) + } + } + } + } + + @Test + fun `should queue messages sent before start`() { + runBlocking { + val message = InitializedNotification() + .toJSON() + + var receivedMessage: JSONRPCMessage? = null + serverTransport.onMessage = { msg -> + receivedMessage = msg + } + + clientTransport.send(message) + serverTransport.start() + assertEquals(message, receivedMessage) + } + } +} \ No newline at end of file diff --git a/src/test/kotlin/client/SseTransportTest.kt b/src/test/kotlin/client/SseTransportTest.kt new file mode 100644 index 0000000..32bb99d --- /dev/null +++ b/src/test/kotlin/client/SseTransportTest.kt @@ -0,0 +1,69 @@ +package client + +import io.ktor.client.* +import io.ktor.client.plugins.sse.* +import io.ktor.server.application.* +import io.ktor.server.cio.* +import io.ktor.server.engine.* +import io.ktor.server.routing.* +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.test.runTest +import mcpSse +import mcpSseTransport +import org.jetbrains.kotlinx.mcp.client.mcpSseTransport +import org.junit.jupiter.api.Test + +private const val PORT = 8080 + +class SseTransportTest : BaseTransportTest() { + @Test + fun `should start then close cleanly`() = runTest { + val server = embeddedServer(CIO, port = PORT) { + install(io.ktor.server.sse.SSE) + routing { + mcpSse() + } + }.start(wait = false) + + val client = HttpClient { + install(SSE) + }.mcpSseTransport { + url { + host = "localhost" + port = PORT + } + } + + testClientOpenClose(client) + + server.stop() + } + + @Test + fun `should read messages`() = runTest { + val clientFinished = CompletableDeferred() + + val server = embeddedServer(CIO, port = PORT) { + install(io.ktor.server.sse.SSE) + routing { + mcpSseTransport { + onMessage = { + send(it) + } + } + } + }.start(wait = false) + + val client = HttpClient { + install(SSE) + }.mcpSseTransport { + url { + host = "localhost" + port = PORT + } + } + + testClientRead(client) + server.stop() + } +} diff --git a/src/test/kotlin/client/StdioClientTransportTest.kt b/src/test/kotlin/client/StdioClientTransportTest.kt new file mode 100644 index 0000000..828858c --- /dev/null +++ b/src/test/kotlin/client/StdioClientTransportTest.kt @@ -0,0 +1,45 @@ +package client + +import kotlinx.coroutines.test.runTest +import org.jetbrains.kotlinx.mcp.client.StdioClientTransport +import org.junit.jupiter.api.Test + +class StdioClientTransportTest : BaseTransportTest() { + @Test + fun `should start then close cleanly`() = runTest { + // Run process "/usr/bin/tee" + val processBuilder = ProcessBuilder("/usr/bin/tee") + val process = processBuilder.start() + + val input = process.inputStream + val output = process.outputStream + + val client = StdioClientTransport( + input = input, + output = output + ) + + testClientOpenClose(client) + + process.destroy() + } + + @Test + fun `should read messages`() = runTest { + val processBuilder = ProcessBuilder("/usr/bin/tee") + val process = processBuilder.start() + + val input = process.inputStream + val output = process.outputStream + + val client = StdioClientTransport( + input = input, + output = output + ) + + testClientRead(client) + + process.waitFor() + process.destroy() + } +} diff --git a/src/test/kotlin/client/TypesTest.kt b/src/test/kotlin/client/TypesTest.kt new file mode 100644 index 0000000..e632178 --- /dev/null +++ b/src/test/kotlin/client/TypesTest.kt @@ -0,0 +1,27 @@ +package client + +import org.jetbrains.kotlinx.mcp.JSONRPCMessage +import org.jetbrains.kotlinx.mcp.Request +import org.junit.jupiter.api.Test +import org.jetbrains.kotlinx.mcp.shared.McpJson + +class TypesTest { + + @Test + fun testRequestResult() { + val message = "{\"result\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{\"tools\":{\"listChanged\":true},\"resources\":{}},\"serverInfo\":{\"name\":\"jetbrains/proxy\",\"version\":\"0.1.0\"}},\"jsonrpc\":\"2.0\",\"id\":1}" + McpJson.decodeFromString(message) + } + + @Test + fun testRequestError() { + val message = "{\"method\":\"initialize\", \"protocolVersion\":\"2024-11-05\",\"capabilities\":{\"experimental\":{},\"sampling\":{}},\"clientInfo\":{\"name\":\"test client\",\"version\":\"1.0\"},\"_meta\":{}}" + McpJson.decodeFromString(message) + } + + @Test + fun testJSONRPCMessage() { + val line = "{\"result\":{\"content\":[{\"type\":\"text\"}],\"isError\":false},\"jsonrpc\":\"2.0\",\"id\":4}" + McpJson.decodeFromString(line) + } +} \ No newline at end of file diff --git a/src/test/kotlin/client/WebSocketTransportTest.kt b/src/test/kotlin/client/WebSocketTransportTest.kt new file mode 100644 index 0000000..8688845 --- /dev/null +++ b/src/test/kotlin/client/WebSocketTransportTest.kt @@ -0,0 +1,49 @@ +package client + +import io.ktor.server.testing.* +import io.ktor.server.websocket.* +import kotlinx.coroutines.CompletableDeferred +import org.jetbrains.kotlinx.mcp.client.mcpWebSocketTransport +import org.junit.jupiter.api.Test +import org.jetbrains.kotlinx.mcp.server.mcpWebSocket +import org.jetbrains.kotlinx.mcp.server.mcpWebSocketTransport + +class WebSocketTransportTest : BaseTransportTest() { + @Test + fun `should start then close cleanly`() = testApplication { + install(WebSockets) + routing { + mcpWebSocket() + } + + val client = createClient { + install(io.ktor.client.plugins.websocket.WebSockets) + }.mcpWebSocketTransport() + + testClientOpenClose(client) + } + + @Test + fun `should read messages`() = testApplication { + val clientFinished = CompletableDeferred() + + install(WebSockets) + routing { + mcpWebSocketTransport { + onMessage = { + send(it) + } + + clientFinished.await() + } + } + + val client = createClient { + install(io.ktor.client.plugins.websocket.WebSockets) + }.mcpWebSocketTransport() + + testClientRead(client) + + clientFinished.complete(Unit) + } +} diff --git a/src/test/kotlin/server/StdioServerTransportTest.kt b/src/test/kotlin/server/StdioServerTransportTest.kt new file mode 100644 index 0000000..ad04226 --- /dev/null +++ b/src/test/kotlin/server/StdioServerTransportTest.kt @@ -0,0 +1,141 @@ +package server + +import org.jetbrains.kotlinx.mcp.InitializedNotification +import org.jetbrains.kotlinx.mcp.JSONRPCMessage +import org.jetbrains.kotlinx.mcp.PingRequest +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.runBlocking +import org.jetbrains.kotlinx.mcp.server.StdioServerTransport +import org.junit.jupiter.api.Assertions.* +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.jetbrains.kotlinx.mcp.shared.ReadBuffer +import org.jetbrains.kotlinx.mcp.shared.serializeMessage +import org.jetbrains.kotlinx.mcp.toJSON +import java.io.* + +class StdioServerTransportTest { + private lateinit var input: PipedInputStream + private lateinit var inputWriter: PipedOutputStream + private lateinit var outputBuffer: ReadBuffer + private lateinit var output: ByteArrayOutputStream + + // We'll store the wrapped streams that meet the constructor requirements + private lateinit var bufferedInput: BufferedInputStream + private lateinit var printOutput: PrintStream + + @BeforeEach + fun setUp() { + // Simulate an input stream that we can push data into using inputWriter. + input = PipedInputStream() + inputWriter = PipedOutputStream(input) + + outputBuffer = ReadBuffer() + + // A custom ByteArrayOutputStream that appends all written data into outputBuffer + output = object : ByteArrayOutputStream() { + override fun write(b: ByteArray, off: Int, len: Int) { + super.write(b, off, len) + outputBuffer.append(b.copyOfRange(off, off + len)) + } + } + + // Wrap input in BufferedInputStream + bufferedInput = BufferedInputStream(input) + + // Wrap output in PrintStream + printOutput = PrintStream(output, true) + } + + @Test + fun `should start then close cleanly`() { + runBlocking { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onError = { error -> + throw error + } + + var didClose = false + server.onClose = { + didClose = true + } + + server.start() + assertFalse(didClose, "Should not have closed yet") + + server.close() + assertTrue(didClose, "Should have closed after calling close()") + } + } + + @Test + fun `should not read until started`() { + runBlocking { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onError = { error -> + throw error + } + + var didRead = false + val readMessage = CompletableDeferred() + + server.onMessage = { message -> + didRead = true + readMessage.complete(message) + } + + val message = PingRequest().toJSON() + + // Push message before the server started + val serialized = serializeMessage(message) + inputWriter.write(serialized) + inputWriter.flush() + + assertFalse(didRead, "Should not have read message before start") + + server.start() + val received = readMessage.await() + assertEquals(message, received) + } + } + + @Test + fun `should read multiple messages`() { + runBlocking { + val server = StdioServerTransport(bufferedInput, printOutput) + server.onError = { error -> + throw error + } + + val messages = listOf( + PingRequest().toJSON(), + InitializedNotification().toJSON(), + ) + + val readMessages = mutableListOf() + val finished = CompletableDeferred() + + server.onMessage = { message -> + readMessages.add(message) + if (message == messages[1]) { + finished.complete(Unit) + } + } + + // Push both messages before starting the server + for (m in messages) { + inputWriter.write(serializeMessage(m)) + } + inputWriter.flush() + + server.start() + finished.await() + + assertEquals(messages, readMessages) + } + } +} + +fun PipedOutputStream.write(s: String) { + write(s.toByteArray()) +} \ No newline at end of file diff --git a/src/test/kotlin/shared/ReadBufferTest.kt b/src/test/kotlin/shared/ReadBufferTest.kt new file mode 100644 index 0000000..5b50c64 --- /dev/null +++ b/src/test/kotlin/shared/ReadBufferTest.kt @@ -0,0 +1,58 @@ +package shared + +import org.jetbrains.kotlinx.mcp.JSONRPCMessage +import org.jetbrains.kotlinx.mcp.JSONRPCNotification +import kotlinx.serialization.encodeToString +import kotlinx.serialization.json.Json +import org.jetbrains.kotlinx.mcp.shared.ReadBuffer +import org.jetbrains.kotlinx.mcp.shared.serializeMessage +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Test +import java.nio.charset.StandardCharsets + +class ReadBufferTest { + private val testMessage: JSONRPCMessage = JSONRPCNotification(method = "foobar") + + private val json = Json { + ignoreUnknownKeys = true + encodeDefaults = true + } + + @Test + fun `should have no messages after initialization`() { + val readBuffer = ReadBuffer() + assertNull(readBuffer.readMessage()) + } + + @Test + fun `should only yield a message after a newline`() { + val readBuffer = ReadBuffer() + + // Append message without newline + val messageBytes = json.encodeToString(testMessage) + .toByteArray(StandardCharsets.UTF_8) + readBuffer.append(messageBytes) + assertNull(readBuffer.readMessage()) + + // Append newline and verify message is now available + readBuffer.append("\n".toByteArray(StandardCharsets.UTF_8)) + assertEquals(testMessage, readBuffer.readMessage()) + assertNull(readBuffer.readMessage()) + } + + @Test + fun `should be reusable after clearing`() { + val readBuffer = ReadBuffer() + + readBuffer.append("foobar".toByteArray(Charsets.UTF_8)) + readBuffer.clear() + assertNull(readBuffer.readMessage()) + + val messageJson = serializeMessage(testMessage) + readBuffer.append(messageJson.toByteArray(Charsets.UTF_8)) + readBuffer.append("\n".toByteArray(Charsets.UTF_8)) + val message = readBuffer.readMessage() + assertEquals(testMessage, message) + } +} diff --git a/src/test/kotlin/sse.ktor.kt b/src/test/kotlin/sse.ktor.kt new file mode 100644 index 0000000..7804ac9 --- /dev/null +++ b/src/test/kotlin/sse.ktor.kt @@ -0,0 +1,162 @@ +import io.ktor.http.HttpStatusCode +import io.ktor.server.application.ApplicationCall +import io.ktor.server.response.respondText +import io.ktor.server.routing.Route +import io.ktor.server.routing.RoutingContext +import io.ktor.server.routing.application +import io.ktor.server.routing.post +import io.ktor.server.sse.ServerSSESession +import io.ktor.server.sse.sse +import io.ktor.util.AttributeKey +import io.ktor.util.Attributes +import kotlinx.coroutines.CompletableDeferred +import org.jetbrains.kotlinx.mcp.Implementation +import org.jetbrains.kotlinx.mcp.ServerCapabilities +import org.jetbrains.kotlinx.mcp.server.SESSION_ID_PARAM +import org.jetbrains.kotlinx.mcp.server.SSEServerTransport +import org.jetbrains.kotlinx.mcp.server.Server +import org.jetbrains.kotlinx.mcp.server.ServerOptions +import org.jetbrains.kotlinx.mcp.shared.IMPLEMENTATION_NAME +import shared.LIB_VERSION + +typealias IncomingHandler = suspend RoutingContext.(forward: suspend () -> Unit) -> Unit + +fun Route.mcpSse( + options: ServerOptions? = null, + incomingPath: String = "", + incomingHandler: (IncomingHandler)? = null, + handler: suspend Server.() -> Unit = {}, +) { + sse { + createMcpServer(this, incomingPath, options, handler) + } + + setupPostRoute(incomingPath, incomingHandler) +} + +fun Route.mcpSse( + path: String, + incomingPath: String = path, + options: ServerOptions? = null, + incomingHandler: (IncomingHandler)? = null, + handler: suspend Server.() -> Unit = {}, +) { + sse(path) { + createMcpServer(this, incomingPath, options, handler) + } + + setupPostRoute(incomingPath, incomingHandler) +} + +fun Route.mcpSseTransport( + incomingPath: String = "", + incomingHandler: (IncomingHandler)? = null, + handler: suspend SSEServerTransport.() -> Unit = {}, +) { + sse { + val transport = createMcpTransport(this, incomingPath) + handler(transport) + transport.start() + transport.close() + } + + setupPostRoute(incomingPath, incomingHandler) +} + +fun Route.mcpSseTransport( + path: String, + incomingPath: String = path, + incomingHandler: (IncomingHandler)? = null, + handler: suspend SSEServerTransport.() -> Unit = {}, +) { + sse(path) { + val transport = createMcpTransport(this, incomingPath) + transport.start() + handler(transport) + transport.close() + } + + setupPostRoute(incomingPath, incomingHandler) +} + +internal val McpServersKey = AttributeKey("mcp-servers") + +private fun String.asAttributeKey() = AttributeKey(this) + +private suspend fun Route.forwardMcpMessage(call: ApplicationCall) { + val sessionId = call.request.queryParameters[SESSION_ID_PARAM] + ?.asAttributeKey() + ?: run { + call.sessionNotFound() + return + } + + application.attributes.getOrNull(McpServersKey) + ?.get(sessionId) + ?.handlePostMessage(call) + ?: call.sessionNotFound() +} + +private suspend fun ApplicationCall.sessionNotFound() { + respondText("Session not found", status = HttpStatusCode.NotFound) +} + +private fun Route.setupPostRoute(incomingPath: String, incomingHandler: IncomingHandler?) { + post(incomingPath) { + if (incomingHandler != null) { + incomingHandler { + forwardMcpMessage(call) + } + } else { + forwardMcpMessage(call) + } + } +} + +private suspend fun Route.createMcpServer( + session: ServerSSESession, + incomingPath: String, + options: ServerOptions?, + handler: suspend Server.() -> Unit = {}, +) { + val transport = createMcpTransport(session, incomingPath) + + val closed = CompletableDeferred() + + 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), + ) + ), + onCloseCallback = { + closed.complete(Unit) + }, + ) + + server.connect(transport) + handler(server) + server.close() +} + +private fun Route.createMcpTransport( + session: ServerSSESession, + incomingPath: String, +): SSEServerTransport { + val transport = SSEServerTransport( + endpoint = incomingPath, + session = session, + ) + + application.attributes + .computeIfAbsent(McpServersKey) { Attributes(concurrent = true) } + .put(transport.sessionId.asAttributeKey(), transport) + + return transport +}