Initial commit

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<JSONRPCMessage> = mutableListOf()
+
+    override var onClose: (() -> Unit)? = null
+    override var onError: ((Throwable) -> Unit)? = null
+    override var onMessage: (suspend ((JSONRPCMessage) -> Unit))? = null
+
+    /**
+     * Creates a pair of linked in-memory transports that can communicate with each other.
+     * One should be passed to a Client and one to a Server.
+     */
+    companion object {
+        fun createLinkedPair(): Pair<InMemoryTransport, InMemoryTransport> {
+            val clientTransport = InMemoryTransport()
+            val serverTransport = InMemoryTransport()
+            clientTransport.otherTransport = serverTransport
+            serverTransport.otherTransport = clientTransport
+            return Pair(clientTransport, serverTransport)
+        }
+    }
+
+    override suspend fun start() {
+        // Process any messages that were queued before start was called
+        while (messageQueue.isNotEmpty()) {
+            messageQueue.removeFirstOrNull()?.let { message ->
+                onMessage?.invoke(message) // todo?
+            }
+        }
+    }
+
+    override suspend fun close() {
+        val other = otherTransport
+        otherTransport = null
+        other?.close()
+        onClose?.invoke()
+    }
+
+    override suspend fun send(message: JSONRPCMessage) {
+        val other = otherTransport ?: throw IllegalStateException("Not connected")
+
+        if (other.onMessage != null) {
+            other.onMessage?.invoke(message) // todo?
+        } else {
+            other.messageQueue.add(message)
+        }
+    }
+}

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<JSONRPCMessage>(
+            PingRequest().toJSON(),
+            InitializedNotification().toJSON()
+        )
+
+        val readMessages = mutableListOf<JSONRPCMessage>()
+        val finished = CompletableDeferred<Unit>()
+
+        client.onMessage = { message ->
+            readMessages.add(message)
+            if (message == messages.last()) {
+                finished.complete(Unit)
+            }
+        }
+
+        client.start()
+
+        for (message in messages) {
+            client.send(message)
+        }
+
+        finished.await()
+
+        assertEquals(messages, readMessages, "Assert messages received")
+
+        client.close()
+    }
+}

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()
+        }
+    }
+
+}

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<IllegalStateException>("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<InitializeRequest>(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<ListResourcesRequest>(Method.Defined.ResourcesList) { request, _ ->
+            ListResourcesResult(resources = emptyList(), nextCursor = null)
+        }
+
+        server.setRequestHandler<ListToolsRequest>(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<IllegalStateException> {
+            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<IllegalStateException> {
+            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<IllegalStateException> {
+            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<Unit>()
+        val defTimeOut = CompletableDeferred<Unit>()
+        server.setRequestHandler<ListResourcesRequest>(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<Unit>()
+        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<ListResourcesRequest>(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<Exception> {
+            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<CreateMessageRequest>(Method.Defined.SamplingCreateMessage) { request, _ ->
+            CreateMessageResult(
+                model = "test-model",
+                role = Role.assistant,
+                content = TextContent(
+                    text = "Test response"
+                )
+            )
+        }
+
+        assertFailsWith<IllegalStateException>("Client does not support roots capability (required for RootsList)") {
+            client.setRequestHandler<ListRootsRequest>(Method.Defined.RootsList) { _, _ -> null }
+        }
+    }
+
+
+}

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<IllegalStateException> {
+                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)
+        }
+    }
+}

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<Unit>()
+
+        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()
+    }
+}

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()
+    }
+}

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<JSONRPCMessage>(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<Request>(message)
+    }
+
+    @Test
+    fun testJSONRPCMessage() {
+        val line = "{\"result\":{\"content\":[{\"type\":\"text\"}],\"isError\":false},\"jsonrpc\":\"2.0\",\"id\":4}"
+        McpJson.decodeFromString<JSONRPCMessage>(line)
+    }
+}

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<Unit>()
+
+        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)
+    }
+}

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<JSONRPCMessage>()
+
+            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<JSONRPCMessage>()
+            val finished = CompletableDeferred<Unit>()
+
+            server.onMessage = { message ->
+                readMessages.add(message)
+                if (message == messages[1]) {
+                    finished.complete(Unit)
+                }
+            }
+
+            // Push both messages before starting the server
+            for (m in messages) {
+                inputWriter.write(serializeMessage(m))
+            }
+            inputWriter.flush()
+
+            server.start()
+            finished.await()
+
+            assertEquals(messages, readMessages)
+        }
+    }
+}
+
+fun PipedOutputStream.write(s: String) {
+    write(s.toByteArray())
+}

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)
+    }
+}

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<Attributes>("mcp-servers")
+
+private fun String.asAttributeKey() = AttributeKey<SSEServerTransport>(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<Unit>()
+
+    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
+}