refactor SseClientTransport (#142)

* refactor `SseClientTransport`

* Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* replace `Exception` with `Throwable` in `SseClientTransport` error handling and try/catch `cancel` in closeResources

---------

Co-authored-by: Leonid Stashevsky <e5l@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SSEClientTransport.kt

@@ -16,6 +16,7 @@ import io.ktor.http.protocolWithAuthority
 import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
 import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
 import io.modelcontextprotocol.kotlin.sdk.shared.McpJson
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
@@ -24,10 +25,11 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.launch
+import kotlinx.serialization.SerializationException
 import kotlin.concurrent.atomics.AtomicBoolean
 import kotlin.concurrent.atomics.ExperimentalAtomicApi
-import kotlin.properties.Delegates
 import kotlin.time.Duration
 
 @Deprecated("Use SseClientTransport instead", ReplaceWith("SseClientTransport"), DeprecationLevel.WARNING)
@@ -44,97 +46,59 @@ public class SseClientTransport(
     private val reconnectionTime: Duration? = null,
     private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
 ) : AbstractTransport() {
-    private val scope by lazy {
-        CoroutineScope(session.coroutineContext + SupervisorJob())
-    }
-
     private val initialized: AtomicBoolean = AtomicBoolean(false)
-    private var session: ClientSSESession by Delegates.notNull()
     private val endpoint = CompletableDeferred<String>()
 
+    private lateinit var session: ClientSSESession
+    private lateinit var scope: CoroutineScope
     private var job: Job? = null
 
-    private val baseUrl by lazy {
+    private val baseUrl: String by lazy {
-        val requestUrl = session.call.request.url.toString()
+        session.call.request.url.let { url ->
-        val url = Url(requestUrl)
+            val path = url.encodedPath
-        var path = url.encodedPath
+            when {
-        if (path.isEmpty()) {
+                path.isEmpty() -> url.protocolWithAuthority
-            url.protocolWithAuthority
+                path.endsWith("/") -> url.protocolWithAuthority + path.removeSuffix("/")
-        } else if (path.endsWith("/")) {
+                else -> url.protocolWithAuthority + path.take(path.lastIndexOf("/"))
-            url.protocolWithAuthority + path.removeSuffix("/")
+            }
-        } else {
-            // the last item is not a directory, so will not be taken into account
-            path = path.substring(0, path.lastIndexOf("/"))
-            url.protocolWithAuthority + path
         }
     }
 
     override suspend fun start() {
-        if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
+        check(initialized.compareAndSet(expectedValue = false, newValue = true)) {
-            error(
+            "SSEClientTransport already started! If using Client class, note that connect() calls start() automatically."
-                "SSEClientTransport already started! " +
-                        "If using Client class, note that connect() calls start() automatically.",
-            )
         }
 
-        session = urlString?.let {
+        try {
-            client.sseSession(
+            session = urlString?.let {
-                urlString = it,
+                client.sseSession(
+                    urlString = it,
+                    reconnectionTime = reconnectionTime,
+                    block = requestBuilder,
+                )
+            } ?: client.sseSession(
                 reconnectionTime = reconnectionTime,
                 block = requestBuilder,
             )
-        } ?: client.sseSession(
+            scope = CoroutineScope(session.coroutineContext + SupervisorJob())
-            reconnectionTime = reconnectionTime,
-            block = requestBuilder,
-        )
 
-        job = scope.launch(CoroutineName("SseMcpClientTransport.collect#${hashCode()}")) {
+            job = scope.launch(CoroutineName("SseMcpClientTransport.connect#${hashCode()}")) {
-            session.incoming.collect { event ->
+                collectMessages()
-                when (event.event) {
-                    "error" -> {
-                        val e = IllegalStateException("SSE error: ${event.data}")
-                        _onError(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/${if (eventData.startsWith("/")) eventData.substring(1) else eventData}")
-                            endpoint.complete(maybeEndpoint.toString())
-                        } catch (e: Exception) {
-                            _onError(e)
-                            close()
-                            error(e)
-                        }
-                    }
-
-                    else -> {
-                        try {
-                            val message = McpJson.decodeFromString<JSONRPCMessage>(event.data ?: "")
-                            _onMessage(message)
-                        } catch (e: Exception) {
-                            _onError(e)
-                        }
-                    }
-                }
             }
-        }
 
-        endpoint.await()
+            endpoint.await()
+        } catch (e: Exception) {
+            closeResources()
+            initialized.store(false)
+            throw e
+        }
     }
 
     @OptIn(ExperimentalCoroutinesApi::class)
     override suspend fun send(message: JSONRPCMessage) {
-        if (!endpoint.isCompleted) {
+        check(initialized.load()) { "SseClientTransport is not initialized!" }
-            error("Not connected")
+        check(job?.isActive == true) { "SseClientTransport is closed!" }
-        }
+        check(endpoint.isCompleted) { "Not connected!" }
 
         try {
             val response = client.post(endpoint.getCompleted()) {
@@ -147,19 +111,80 @@ public class SseClientTransport(
                 val text = response.bodyAsText()
                 error("Error POSTing to endpoint (HTTP ${response.status}): $text")
             }
-        } catch (e: Exception) {
+        } catch (e: Throwable) {
             _onError(e)
             throw e
         }
     }
 
     override suspend fun close() {
-        if (!initialized.load()) {
+        check(initialized.load()) { "SseClientTransport is not initialized!" }
-            error("SSEClientTransport is not initialized!")
+        closeResources()
+    }
+
+    private suspend fun CoroutineScope.collectMessages() {
+        try {
+            session.incoming.collect { event ->
+                ensureActive()
+
+                when (event.event) {
+                    "error" -> {
+                        val error = IllegalStateException("SSE error: ${event.data}")
+                        _onError(error)
+                        throw error
+                    }
+
+                    "open" -> {
+                        // The connection is open, but we need to wait for the endpoint to be received.
+                    }
+
+                    "endpoint" -> handleEndpoint(event.data.orEmpty())
+                    else -> handleMessage(event.data.orEmpty())
+                }
+            }
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Throwable) {
+            _onError(e)
+            throw e
+        } finally {
+            closeResources()
+        }
+    }
+
+    private fun handleEndpoint(eventData: String) {
+        try {
+            val path = if (eventData.startsWith("/")) eventData.substring(1) else eventData
+            val endpointUrl = Url("$baseUrl/$path")
+            endpoint.complete(endpointUrl.toString())
+        } catch (e: Throwable) {
+            _onError(e)
+            endpoint.completeExceptionally(e)
+            throw e
+        }
+    }
+
+    private suspend fun handleMessage(data: String) {
+        try {
+            val message = McpJson.decodeFromString<JSONRPCMessage>(data)
+            _onMessage(message)
+        } catch (e: SerializationException) {
+            _onError(e)
+        }
+    }
+
+    private suspend fun closeResources() {
+        if (!initialized.compareAndSet(expectedValue = true, newValue = false)) return
+
+        job?.cancelAndJoin()
+        try {
+            if (::session.isInitialized) session.cancel()
+            if (::scope.isInitialized) scope.cancel()
+            endpoint.cancel()
+        } catch (e: Throwable) {
+            _onError(e)
         }
 
-        session.cancel()
         _onClose()
-        job?.cancelAndJoin()
     }
 }

src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/client/SseTransportTest.kt

@@ -4,6 +4,7 @@ import io.ktor.client.HttpClient
 import io.ktor.client.plugins.sse.SSE
 import io.ktor.server.application.install
 import io.ktor.server.cio.CIO
+import io.ktor.server.engine.EmbeddedServer
 import io.ktor.server.engine.embeddedServer
 import io.ktor.server.routing.post
 import io.ktor.server.routing.route
@@ -17,10 +18,12 @@ import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 
 class SseTransportTest : BaseTransportTest() {
+
+    private suspend fun EmbeddedServer<*, *>.actualPort() = engine.resolvedConnectors().first().port
+
     @Test
     fun `should start then close cleanly`() = runTest {
-        val port = 8080
+        val server = embeddedServer(CIO, port = 0) {
-        val server = embeddedServer(CIO, port = port) {
             install(io.ktor.server.sse.SSE)
             val transports = ConcurrentMap<String, SseServerTransport>()
             routing {
@@ -34,24 +37,27 @@ class SseTransportTest : BaseTransportTest() {
             }
         }.startSuspend(wait = false)
 
+        val actualPort = server.actualPort()
+
         val client = HttpClient {
             install(SSE)
         }.mcpSseTransport {
             url {
                 host = "localhost"
-                this.port = port
+                this.port = actualPort
             }
         }
 
-        testClientOpenClose(client)
+        try {
-
+            testClientOpenClose(client)
-        server.stopSuspend()
+        } finally {
+            server.stopSuspend()
+        }
     }
 
     @Test
     fun `should read messages`() = runTest {
-        val port = 3003
+        val server = embeddedServer(CIO, port = 0) {
-        val server = embeddedServer(CIO, port = port) {
             install(io.ktor.server.sse.SSE)
             val transports = ConcurrentMap<String, SseServerTransport>()
             routing {
@@ -71,23 +77,27 @@ class SseTransportTest : BaseTransportTest() {
             }
         }.startSuspend(wait = false)
 
+        val actualPort = server.actualPort()
+
         val client = HttpClient {
             install(SSE)
         }.mcpSseTransport {
             url {
                 host = "localhost"
-                this.port = port
+                this.port = actualPort
             }
         }
 
-        testClientRead(client)
+        try {
-        server.stopSuspend()
+            testClientRead(client)
+        } finally {
+            server.stopSuspend()
+        }
     }
 
     @Test
     fun `test sse path not root path`() = runTest {
-        val port = 3007
+        val server = embeddedServer(CIO, port = 0) {
-        val server = embeddedServer(CIO, port = port) {
             install(io.ktor.server.sse.SSE)
             val transports = ConcurrentMap<String, SseServerTransport>()
             routing {
@@ -109,17 +119,22 @@ class SseTransportTest : BaseTransportTest() {
             }
         }.startSuspend(wait = false)
 
+        val actualPort = server.actualPort()
+
         val client = HttpClient {
             install(SSE)
         }.mcpSseTransport {
             url {
                 host = "localhost"
-                this.port = port
+                this.port = actualPort
                 pathSegments = listOf("sse")
             }
         }
 
-        testClientRead(client)
+        try {
-        server.stopSuspend()
+            testClientRead(client)
+        } finally {
+            server.stopSuspend()
+        }
     }
 }

src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/SseIntegrationTest.kt

@@ -26,11 +26,15 @@ class SseIntegrationTest {
     @Test
     fun `client should be able to connect to sse server`() = runTest {
         val serverEngine = initServer()
+        var client: Client? = null
         try {
             withContext(Dispatchers.Default) {
-                assertDoesNotThrow { initClient() }
+                assertDoesNotThrow { client = initClient() }
             }
+        } catch (e: Exception) {
+            fail("Failed to connect client: $e")
         } finally {
+            client?.close()
             // Make sure to stop the server
             serverEngine.stopSuspend(1000, 2000)
         }