## refactor: Introduce AbstractClientTransport and simplify client transport lifecycle This PR unifies state management logic across client transports, except for the websocket transport, as it needs more refactoring first and is not the primary transport. - Introduce `AbstractClientTransport` as base class for other client transports. - Migrate to `ClientTransportState` states with state machine, streamline state management. - Ensure proper cleanup and error handling consistency across transport implementations. - Add unit tests to validate `AbstractClientTransport` behavior and state transitions. - Update related documentation and API contract. - Make sure `onClose` callback is called at most once - fix [StreamableHttpClientTransportTest.kt](https://github.com/modelcontextprotocol/kotlin-sdk/pull/516/changes#diff-b9fe8891d08a492448a7ed2dc782151c6a915ae1403a4d70c2034c6062c84cb5) to throw `McpException` on connect when received invalid Json ## Motivation and Context Unify lifecycle handling and eliminate duplication in client Transport implementations. ## How Has This Been Tested? Unit test, regression tests ## Breaking Changes - `onClose` callback is now called only once - send(...) will throw McpException with message "Transport is not ready" when `ClientTransportState != Operational` ## Types of changes - [ ] Bug fix (non-breaking change which fixes an issue) - [x] Refactoring - [ ] New feature (non-breaking change which adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update ## Checklist <!-- Go over all the following points, and put an `x` in all the boxes that apply. --> - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [ ] New and existing tests pass locally - [ ] I have added appropriate error handling - [ ] I have added or updated documentation as needed ## Additional context See #517
This commit is contained in:
Konstantin Pavlov
GitHub
parent
99908beb2a
commit
6120b3109e
24 changed files with 1057 additions and 191 deletions
|
|
@ -417,6 +417,6 @@ private class HttpServerTransport(private val sessionId: String) : AbstractTrans
|
|||
logger.debug { "Closing transport for session $sessionId" }
|
||||
messageQueue.close()
|
||||
pendingResponses.clear()
|
||||
_onClose.invoke()
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ class InMemoryTransport : AbstractTransport() {
|
|||
val other = otherTransport
|
||||
otherTransport = null
|
||||
other?.close()
|
||||
_onClose.invoke()
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
|
|
|
|||
|
|
@ -455,7 +455,7 @@ class HttpServerTransport(private val sessionId: String) : AbstractTransport() {
|
|||
override suspend fun close() {
|
||||
logger.debug { "Closing HTTP server transport for session: $sessionId" }
|
||||
messageQueue.close()
|
||||
_onClose.invoke()
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1108,7 +1108,6 @@
|
|||
"resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz",
|
||||
"integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==",
|
||||
"dev": true,
|
||||
"hasInstallScript": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"os": [
|
||||
|
|
|
|||
|
|
@ -62,24 +62,18 @@ public final class io/modelcontextprotocol/kotlin/sdk/client/KtorClientKt {
|
|||
public static synthetic fun mcpSseTransport-5_5nbZA$default (Lio/ktor/client/HttpClient;Ljava/lang/String;Lkotlin/time/Duration;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/client/SseClientTransport;
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/SseClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport {
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/SseClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractClientTransport {
|
||||
public synthetic fun <init> (Lio/ktor/client/HttpClient;Ljava/lang/String;Lkotlin/time/Duration;Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public synthetic fun <init> (Lio/ktor/client/HttpClient;Ljava/lang/String;Lkotlin/time/Duration;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public fun start (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport {
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractClientTransport {
|
||||
public static final field BUFFER_SIZE J
|
||||
public fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;)V
|
||||
public fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;Lkotlinx/io/Source;)V
|
||||
public fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;Lkotlinx/io/Source;Lkotlinx/coroutines/channels/Channel;)V
|
||||
public fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;Lkotlinx/io/Source;Lkotlinx/coroutines/channels/Channel;Lkotlin/jvm/functions/Function1;)V
|
||||
public synthetic fun <init> (Lkotlinx/io/Source;Lkotlinx/io/Sink;Lkotlinx/io/Source;Lkotlinx/coroutines/channels/Channel;Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public fun start (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport$StderrSeverity : java/lang/Enum {
|
||||
|
|
@ -93,17 +87,14 @@ public final class io/modelcontextprotocol/kotlin/sdk/client/StdioClientTranspor
|
|||
public static fun values ()[Lio/modelcontextprotocol/kotlin/sdk/client/StdioClientTransport$StderrSeverity;
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport {
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractClientTransport {
|
||||
public synthetic fun <init> (Lio/ktor/client/HttpClient;Ljava/lang/String;Lkotlin/time/Duration;Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public synthetic fun <init> (Lio/ktor/client/HttpClient;Ljava/lang/String;Lkotlin/time/Duration;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public final fun getProtocolVersion ()Ljava/lang/String;
|
||||
public final fun getSessionId ()Ljava/lang/String;
|
||||
public fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public final fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Ljava/lang/String;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public static synthetic fun send$default (Lio/modelcontextprotocol/kotlin/sdk/client/StreamableHttpClientTransport;Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Ljava/lang/String;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
|
||||
public final fun setProtocolVersion (Ljava/lang/String;)V
|
||||
public fun start (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public final fun terminateSession (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,170 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.client
|
||||
|
||||
import io.github.oshai.kotlinlogging.KLogger
|
||||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.client.plugins.sse.ClientSSESession
|
||||
import io.ktor.client.plugins.sse.sseSession
|
||||
import io.ktor.client.request.HttpRequestBuilder
|
||||
import io.ktor.client.request.post
|
||||
import io.ktor.client.request.setBody
|
||||
import io.ktor.client.statement.bodyAsText
|
||||
import io.ktor.http.ContentType
|
||||
import io.ktor.http.HttpHeaders
|
||||
import io.ktor.http.Url
|
||||
import io.ktor.http.append
|
||||
import io.ktor.http.isSuccess
|
||||
import io.ktor.http.protocolWithAuthority
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractClientTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpJson
|
||||
import kotlinx.coroutines.CancellationException
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.CoroutineName
|
||||
import kotlinx.coroutines.CoroutineScope
|
||||
import kotlinx.coroutines.ExperimentalCoroutinesApi
|
||||
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.ExperimentalAtomicApi
|
||||
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.
|
||||
*/
|
||||
@OptIn(ExperimentalAtomicApi::class)
|
||||
public class SseClientTransport(
|
||||
private val client: HttpClient,
|
||||
private val urlString: String?,
|
||||
private val reconnectionTime: Duration? = null,
|
||||
private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
|
||||
) : AbstractClientTransport() {
|
||||
|
||||
override val logger: KLogger = KotlinLogging.logger {}
|
||||
|
||||
private val endpoint = CompletableDeferred<String>()
|
||||
|
||||
private lateinit var session: ClientSSESession
|
||||
private lateinit var scope: CoroutineScope
|
||||
private var job: Job? = null
|
||||
|
||||
private val baseUrl: String by lazy {
|
||||
session.call.request.url.let { url ->
|
||||
val path = url.encodedPath
|
||||
when {
|
||||
path.isEmpty() -> url.protocolWithAuthority
|
||||
path.endsWith("/") -> url.protocolWithAuthority + path.removeSuffix("/")
|
||||
else -> url.protocolWithAuthority + path.take(path.lastIndexOf("/"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun initialize() {
|
||||
session = urlString?.let {
|
||||
client.sseSession(
|
||||
urlString = it,
|
||||
reconnectionTime = reconnectionTime,
|
||||
block = requestBuilder,
|
||||
)
|
||||
} ?: client.sseSession(
|
||||
reconnectionTime = reconnectionTime,
|
||||
block = requestBuilder,
|
||||
)
|
||||
scope = CoroutineScope(session.coroutineContext + SupervisorJob())
|
||||
|
||||
job = scope.launch(CoroutineName("SseMcpClientTransport.connect#${hashCode()}")) {
|
||||
collectMessages()
|
||||
}
|
||||
|
||||
endpoint.await()
|
||||
}
|
||||
|
||||
@OptIn(ExperimentalCoroutinesApi::class)
|
||||
override suspend fun performSend(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
check(job?.isActive == true) { "SseClientTransport is closed!" }
|
||||
check(endpoint.isCompleted) { "Not connected!" }
|
||||
|
||||
val response = client.post(endpoint.getCompleted()) {
|
||||
requestBuilder()
|
||||
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")
|
||||
}
|
||||
|
||||
logger.debug { "Client successfully sent message via SSE $endpoint" }
|
||||
}
|
||||
|
||||
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())
|
||||
logger.debug { "Client connected to endpoint: $endpointUrl" }
|
||||
} 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)
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun closeResources() {
|
||||
job?.cancelAndJoin()
|
||||
try {
|
||||
if (::session.isInitialized) session.cancel()
|
||||
if (::scope.isInitialized) scope.cancel()
|
||||
endpoint.cancel()
|
||||
} catch (e: Throwable) {
|
||||
_onError(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.client
|
||||
|
||||
import io.github.oshai.kotlinlogging.KLogger
|
||||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport.StderrSeverity.DEBUG
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport.StderrSeverity.FATAL
|
||||
|
|
@ -7,7 +8,7 @@ import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport.StderrSeve
|
|||
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport.StderrSeverity.INFO
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.StdioClientTransport.StderrSeverity.WARNING
|
||||
import io.modelcontextprotocol.kotlin.sdk.internal.IODispatcher
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractClientTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.ReadBuffer
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.serializeMessage
|
||||
|
|
@ -39,7 +40,6 @@ import kotlinx.io.buffered
|
|||
import kotlinx.io.readByteArray
|
||||
import kotlinx.io.writeString
|
||||
import kotlinx.serialization.SerializationException
|
||||
import kotlin.concurrent.atomics.AtomicBoolean
|
||||
import kotlin.concurrent.atomics.ExperimentalAtomicApi
|
||||
import kotlin.coroutines.CoroutineContext
|
||||
import kotlin.jvm.JvmOverloads
|
||||
|
|
@ -88,7 +88,9 @@ public class StdioClientTransport @JvmOverloads public constructor(
|
|||
private val error: Source? = null,
|
||||
private val sendChannel: Channel<JSONRPCMessage> = Channel(Channel.BUFFERED),
|
||||
private val classifyStderr: (String) -> StderrSeverity = { DEBUG },
|
||||
) : AbstractTransport() {
|
||||
) : AbstractClientTransport() {
|
||||
|
||||
override val logger: KLogger = KotlinLogging.logger {}
|
||||
|
||||
private companion object {
|
||||
/**
|
||||
|
|
@ -96,8 +98,6 @@ public class StdioClientTransport @JvmOverloads public constructor(
|
|||
* 8KB is optimal for most systems (matches default page size).
|
||||
*/
|
||||
const val BUFFER_SIZE = 8 * 1024L
|
||||
|
||||
private val logger = KotlinLogging.logger {}
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -113,13 +113,8 @@ public class StdioClientTransport @JvmOverloads public constructor(
|
|||
|
||||
private val ioCoroutineContext: CoroutineContext = IODispatcher
|
||||
private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
|
||||
private val initialized = AtomicBoolean(false)
|
||||
private val onCloseCalled = AtomicBoolean(false)
|
||||
|
||||
override suspend fun start() {
|
||||
if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
|
||||
error("StdioClientTransport already started!")
|
||||
}
|
||||
override suspend fun initialize() {
|
||||
logger.debug { "Starting StdioClientTransport..." }
|
||||
|
||||
// Producers run on IODispatcher for I/O
|
||||
|
|
@ -224,24 +219,12 @@ public class StdioClientTransport @JvmOverloads public constructor(
|
|||
// Wait for write job to complete before closing, matching old implementation
|
||||
writeJob?.cancelAndJoin()
|
||||
logger.debug { "Transport coroutine completed, calling onClose" }
|
||||
callOnCloseOnce()
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
if (!initialized.load()) {
|
||||
throw McpException(
|
||||
code = CONNECTION_CLOSED,
|
||||
message = "Transport is not started",
|
||||
)
|
||||
}
|
||||
if (onCloseCalled.load()) {
|
||||
throw McpException(
|
||||
code = CONNECTION_CLOSED,
|
||||
message = "Transport is closed",
|
||||
)
|
||||
}
|
||||
override suspend fun performSend(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
@Suppress("SwallowedException")
|
||||
try {
|
||||
sendChannel.send(message)
|
||||
|
|
@ -254,33 +237,14 @@ public class StdioClientTransport @JvmOverloads public constructor(
|
|||
message = "Transport is closed",
|
||||
cause = e,
|
||||
)
|
||||
} catch (e: McpException) {
|
||||
logger.debug(e) { "Error while sending message: ${e.message}" }
|
||||
throw e
|
||||
} catch (e: Throwable) {
|
||||
logger.error(e) { "Error while sending message: ${e.message}" }
|
||||
throw McpException(
|
||||
code = INTERNAL_ERROR,
|
||||
message = "Error while sending message: ${e.message}",
|
||||
cause = e,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
if (!initialized.compareAndSet(expectedValue = true, newValue = false)) {
|
||||
return // Already closed
|
||||
}
|
||||
override suspend fun closeResources() {
|
||||
scope.stopProcessing("Closed")
|
||||
scope.coroutineContext[Job]?.join() // Wait for all coroutines to complete
|
||||
}
|
||||
|
||||
private fun callOnCloseOnce() {
|
||||
if (onCloseCalled.compareAndSet(expectedValue = false, newValue = true)) {
|
||||
runCatching { _onClose() }
|
||||
}
|
||||
}
|
||||
|
||||
private fun sendOutboundMessage(message: JSONRPCMessage, sink: Sink, mainScope: CoroutineScope) {
|
||||
try {
|
||||
val json = serializeMessage(message)
|
||||
|
|
@ -308,7 +272,7 @@ public class StdioClientTransport @JvmOverloads public constructor(
|
|||
|
||||
private fun CoroutineScope.stopProcessing(reason: String, cause: Throwable? = null) {
|
||||
sendChannel.close() // Stop accepting new messages
|
||||
callOnCloseOnce()
|
||||
invokeOnCloseCallback()
|
||||
cancel(reason, cause) // cancel current coroutine context
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.client
|
||||
|
||||
import io.github.oshai.kotlinlogging.KLogger
|
||||
import io.github.oshai.kotlinlogging.KotlinLogging
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.client.plugins.sse.ClientSSESession
|
||||
|
|
@ -21,15 +22,13 @@ import io.ktor.http.HttpStatusCode
|
|||
import io.ktor.http.contentType
|
||||
import io.ktor.http.isSuccess
|
||||
import io.ktor.utils.io.readUTF8Line
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractClientTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpException
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpJson
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode.CONNECTION_CLOSED
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import kotlinx.coroutines.CancellationException
|
||||
import kotlinx.coroutines.CoroutineName
|
||||
|
|
@ -40,7 +39,6 @@ import kotlinx.coroutines.SupervisorJob
|
|||
import kotlinx.coroutines.cancel
|
||||
import kotlinx.coroutines.cancelAndJoin
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlin.concurrent.atomics.AtomicBoolean
|
||||
import kotlin.concurrent.atomics.ExperimentalAtomicApi
|
||||
import kotlin.time.Duration
|
||||
|
||||
|
|
@ -65,18 +63,14 @@ public class StreamableHttpClientTransport(
|
|||
private val url: String,
|
||||
private val reconnectionTime: Duration? = null,
|
||||
private val requestBuilder: HttpRequestBuilder.() -> Unit = {},
|
||||
) : AbstractTransport() {
|
||||
) : AbstractClientTransport() {
|
||||
|
||||
private companion object {
|
||||
private val logger = KotlinLogging.logger {}
|
||||
}
|
||||
override val logger: KLogger = KotlinLogging.logger {}
|
||||
|
||||
public var sessionId: String? = null
|
||||
private set
|
||||
public var protocolVersion: String? = null
|
||||
|
||||
private val initialized: AtomicBoolean = AtomicBoolean(false)
|
||||
|
||||
private var sseSession: ClientSSESession? = null
|
||||
private var sseJob: Job? = null
|
||||
|
||||
|
|
@ -84,42 +78,21 @@ public class StreamableHttpClientTransport(
|
|||
|
||||
private var lastEventId: String? = null
|
||||
|
||||
override suspend fun start() {
|
||||
if (!initialized.compareAndSet(expectedValue = false, newValue = true)) {
|
||||
error("StreamableHttpClientTransport already started!")
|
||||
}
|
||||
logger.debug { "Client transport starting..." }
|
||||
override suspend fun initialize() {
|
||||
logger.debug { "Client transport is starting..." }
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a single message with optional resumption support
|
||||
*/
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
send(message, options?.resumptionToken, options?.onResumptionToken)
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends one or more messages with optional resumption support.
|
||||
* This is the main send method that matches the TypeScript implementation.
|
||||
*/
|
||||
public suspend fun send(
|
||||
message: JSONRPCMessage,
|
||||
resumptionToken: String?,
|
||||
onResumptionToken: ((String) -> Unit)? = null,
|
||||
) {
|
||||
if (!initialized.load()) {
|
||||
throw McpException(
|
||||
code = CONNECTION_CLOSED,
|
||||
message = "Transport is not started",
|
||||
)
|
||||
}
|
||||
override suspend fun performSend(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
logger.debug { "Client sending message via POST to $url: ${McpJson.encodeToString(message)}" }
|
||||
|
||||
// If we have a resumption token, reconnect the SSE stream with it
|
||||
resumptionToken?.let { token ->
|
||||
options?.resumptionToken?.let { token ->
|
||||
startSseSession(
|
||||
resumptionToken = token,
|
||||
onResumptionToken = onResumptionToken,
|
||||
onResumptionToken = options.onResumptionToken,
|
||||
replayMessageId = if (message is JSONRPCRequest) message.id else null,
|
||||
)
|
||||
return
|
||||
|
|
@ -138,7 +111,7 @@ public class StreamableHttpClientTransport(
|
|||
|
||||
if (response.status == HttpStatusCode.Accepted) {
|
||||
if (message is JSONRPCNotification && message.method == "notifications/initialized") {
|
||||
startSseSession(onResumptionToken = onResumptionToken)
|
||||
startSseSession(onResumptionToken = options?.onResumptionToken)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
|
@ -161,7 +134,7 @@ public class StreamableHttpClientTransport(
|
|||
|
||||
ContentType.Text.EventStream -> handleInlineSse(
|
||||
response,
|
||||
onResumptionToken = onResumptionToken,
|
||||
onResumptionToken = options?.onResumptionToken,
|
||||
replayMessageId = if (message is JSONRPCRequest) message.id else null,
|
||||
)
|
||||
|
||||
|
|
@ -177,23 +150,31 @@ public class StreamableHttpClientTransport(
|
|||
}
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
if (!initialized.load()) return // Already closed or never started
|
||||
/**
|
||||
* Sends one or more messages with optional resumption support.
|
||||
* This is the main send method that matches the TypeScript implementation.
|
||||
*/
|
||||
public suspend fun send(
|
||||
message: JSONRPCMessage,
|
||||
resumptionToken: String?,
|
||||
onResumptionToken: ((String) -> Unit)? = null,
|
||||
): Unit = send(
|
||||
message = message,
|
||||
options = TransportSendOptions(
|
||||
resumptionToken = resumptionToken,
|
||||
onResumptionToken = onResumptionToken,
|
||||
),
|
||||
)
|
||||
|
||||
override suspend fun closeResources() {
|
||||
logger.debug { "Client transport closing." }
|
||||
|
||||
try {
|
||||
// Try to terminate session if we have one
|
||||
terminateSession()
|
||||
// Try to terminate session if we have one
|
||||
terminateSession()
|
||||
|
||||
sseSession?.cancel()
|
||||
sseJob?.cancelAndJoin()
|
||||
scope.cancel()
|
||||
} catch (_: Exception) {
|
||||
// Ignore errors during cleanup
|
||||
} finally {
|
||||
initialized.store(false)
|
||||
_onClose()
|
||||
}
|
||||
sseSession?.cancel()
|
||||
sseJob?.cancelAndJoin()
|
||||
scope.cancel()
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ abstract class AbstractClientTransportLifecycleTest<T : AbstractTransport> {
|
|||
val exception = shouldThrow<IllegalStateException> {
|
||||
transport.start()
|
||||
}
|
||||
exception.message shouldContain "already started"
|
||||
exception.message shouldContain "Can't change state: expected transport state"
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -53,8 +53,8 @@ abstract class AbstractClientTransportLifecycleTest<T : AbstractTransport> {
|
|||
val exception = shouldThrow<McpException> {
|
||||
transport.send(PingRequest().toJSON())
|
||||
}
|
||||
exception.message shouldContain "not started"
|
||||
exception.code shouldBe CONNECTION_CLOSED
|
||||
exception.message shouldContain "Transport is not ready"
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.client.streamable.http
|
||||
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.client.engine.mock.MockEngine
|
||||
import io.ktor.client.engine.mock.MockRequestHandler
|
||||
|
|
@ -18,7 +19,9 @@ import io.modelcontextprotocol.kotlin.sdk.types.Implementation
|
|||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpException
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpJson
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.coroutines.delay
|
||||
import kotlinx.coroutines.test.runTest
|
||||
|
|
@ -94,7 +97,11 @@ class StreamableHttpClientTransportTest {
|
|||
)
|
||||
|
||||
val transport = createTransport { request ->
|
||||
when (val msg = McpJson.decodeFromString<JSONRPCMessage>((request.body as TextContent).text)) {
|
||||
when (
|
||||
val msg = McpJson.decodeFromString<JSONRPCMessage>(
|
||||
(request.body as TextContent).text,
|
||||
)
|
||||
) {
|
||||
is JSONRPCRequest if msg.method == "initialize" -> respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.OK,
|
||||
|
|
@ -407,7 +414,7 @@ class StreamableHttpClientTransportTest {
|
|||
|
||||
try {
|
||||
// Real time-keeping is needed; otherwise Protocol will always throw TimeoutCancellationException in tests
|
||||
assertFailsWith<IllegalStateException>(
|
||||
val mcpException = assertFailsWith<McpException>(
|
||||
message = "Expected client.connect to fail on invalid JSON response",
|
||||
) {
|
||||
withContext(Dispatchers.Default.limitedParallelism(1)) {
|
||||
|
|
@ -416,6 +423,7 @@ class StreamableHttpClientTransportTest {
|
|||
}
|
||||
}
|
||||
}
|
||||
mcpException.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR
|
||||
} finally {
|
||||
transport.close()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -166,18 +166,12 @@ class StdioClientTransportErrorHandlingTest {
|
|||
true,
|
||||
ErrorCode.INTERNAL_ERROR,
|
||||
),
|
||||
Arguments.of(
|
||||
OutOfMemoryError(),
|
||||
true,
|
||||
ErrorCode.INTERNAL_ERROR,
|
||||
),
|
||||
|
||||
)
|
||||
}
|
||||
|
||||
@ParameterizedTest
|
||||
@MethodSource("exceptions")
|
||||
fun `Send should handle exceptions`(throwable: Throwable, shouldWrap: Boolean, expectedCode: Int?) = runTest {
|
||||
fun `Send should handle exceptions`(throwable: Exception, shouldWrap: Boolean, expectedCode: Int?) = runTest {
|
||||
val sendChannel: Channel<JSONRPCMessage> = mockk(relaxed = true)
|
||||
|
||||
// Create stdin pipe that stays open to prevent transport from closing
|
||||
|
|
|
|||
|
|
@ -139,6 +139,9 @@ public final class io/modelcontextprotocol/kotlin/sdk/InitializedNotification {
|
|||
public static synthetic fun Params$default (Lio/modelcontextprotocol/kotlin/sdk/InitializedNotification;Lkotlinx/serialization/json/JsonObject;ILjava/lang/Object;)Lio/modelcontextprotocol/kotlin/sdk/types/BaseNotificationParams;
|
||||
}
|
||||
|
||||
public abstract interface annotation class io/modelcontextprotocol/kotlin/sdk/InternalMcpApi : java/lang/annotation/Annotation {
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/LibVersionKt {
|
||||
public static final field LIB_VERSION Ljava/lang/String;
|
||||
}
|
||||
|
|
@ -408,16 +411,46 @@ public final class io/modelcontextprotocol/kotlin/sdk/internal/Utils_jvmKt {
|
|||
public static final fun getIODispatcher ()Lkotlinx/coroutines/CoroutineDispatcher;
|
||||
}
|
||||
|
||||
public abstract class io/modelcontextprotocol/kotlin/sdk/shared/AbstractClientTransport : io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport {
|
||||
public fun <init> ()V
|
||||
protected final fun checkState (Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;Lkotlin/jvm/functions/Function1;)V
|
||||
public static synthetic fun checkState$default (Lio/modelcontextprotocol/kotlin/sdk/shared/AbstractClientTransport;Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
|
||||
public fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
protected abstract fun closeResources (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
protected abstract fun getLogger ()Lio/github/oshai/kotlinlogging/KLogger;
|
||||
protected final fun getState ()Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
protected abstract fun initialize (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
protected abstract fun performSend (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public static synthetic fun performSend$default (Lio/modelcontextprotocol/kotlin/sdk/shared/AbstractClientTransport;Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
|
||||
public fun send (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;Lio/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public fun start (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
protected final fun stateTransition (Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;)V
|
||||
protected final fun updateState (Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;)V
|
||||
}
|
||||
|
||||
public abstract class io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport : io/modelcontextprotocol/kotlin/sdk/shared/Transport {
|
||||
public fun <init> ()V
|
||||
protected final fun get_onClose ()Lkotlin/jvm/functions/Function0;
|
||||
protected final fun get_onError ()Lkotlin/jvm/functions/Function1;
|
||||
protected final fun get_onMessage ()Lkotlin/jvm/functions/Function2;
|
||||
protected final fun invokeOnCloseCallback ()V
|
||||
public fun onClose (Lkotlin/jvm/functions/Function0;)V
|
||||
public fun onError (Lkotlin/jvm/functions/Function1;)V
|
||||
public fun onMessage (Lkotlin/jvm/functions/Function2;)V
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState : java/lang/Enum {
|
||||
public static final field InitializationFailed Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static final field Initializing Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static final field New Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static final field Operational Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static final field ShutdownFailed Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static final field ShuttingDown Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static final field Stopped Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static fun getEntries ()Lkotlin/enums/EnumEntries;
|
||||
public static fun valueOf (Ljava/lang/String;)Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
public static fun values ()[Lio/modelcontextprotocol/kotlin/sdk/shared/ClientTransportState;
|
||||
}
|
||||
|
||||
public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol {
|
||||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions;)V
|
||||
protected abstract fun assertCapabilityForMethod (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V
|
||||
|
|
|
|||
|
|
@ -130,6 +130,8 @@ kotlin {
|
|||
|
||||
jvmTest {
|
||||
dependencies {
|
||||
implementation(libs.junit.jupiter.params)
|
||||
implementation(libs.mockk)
|
||||
runtimeOnly(libs.slf4j.simple)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,49 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk
|
||||
|
||||
import kotlin.annotation.AnnotationTarget.ANNOTATION_CLASS
|
||||
import kotlin.annotation.AnnotationTarget.CLASS
|
||||
import kotlin.annotation.AnnotationTarget.CONSTRUCTOR
|
||||
import kotlin.annotation.AnnotationTarget.FIELD
|
||||
import kotlin.annotation.AnnotationTarget.FUNCTION
|
||||
import kotlin.annotation.AnnotationTarget.LOCAL_VARIABLE
|
||||
import kotlin.annotation.AnnotationTarget.PROPERTY
|
||||
import kotlin.annotation.AnnotationTarget.PROPERTY_GETTER
|
||||
import kotlin.annotation.AnnotationTarget.PROPERTY_SETTER
|
||||
import kotlin.annotation.AnnotationTarget.TYPEALIAS
|
||||
import kotlin.annotation.AnnotationTarget.VALUE_PARAMETER
|
||||
|
||||
/**
|
||||
* Denotes that the annotated API is internal to the MCP SDK and its usage is restricted to components
|
||||
* within the SDK. It is not intended to be used in general application code or by external consumers.
|
||||
*
|
||||
* This annotation is used to communicate and enforce the limited visibility of certain APIs to prevent
|
||||
* accidental usage outside their intended context.
|
||||
*
|
||||
* Applying this annotation requires consumers to explicitly opt-in, acknowledging the warning level
|
||||
* associated with its use.
|
||||
*
|
||||
* This annotation can be used on classes, annotations, properties, fields, local variables, value
|
||||
* parameters, constructors, functions, property accessors, and type aliases.
|
||||
*
|
||||
* The annotation is retained in the binary and must be documented.
|
||||
*/
|
||||
@RequiresOptIn(
|
||||
message = "This API is internal, and its usage is restricted to internal components of the MCP SDK.",
|
||||
level = RequiresOptIn.Level.WARNING,
|
||||
)
|
||||
@MustBeDocumented
|
||||
@Target(
|
||||
CLASS,
|
||||
ANNOTATION_CLASS,
|
||||
PROPERTY,
|
||||
FIELD,
|
||||
LOCAL_VARIABLE,
|
||||
VALUE_PARAMETER,
|
||||
CONSTRUCTOR,
|
||||
FUNCTION,
|
||||
PROPERTY_GETTER,
|
||||
PROPERTY_SETTER,
|
||||
TYPEALIAS,
|
||||
)
|
||||
@Retention(AnnotationRetention.BINARY)
|
||||
public annotation class InternalMcpApi
|
||||
|
|
@ -0,0 +1,274 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
import io.github.oshai.kotlinlogging.KLogger
|
||||
import io.modelcontextprotocol.kotlin.sdk.InternalMcpApi
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpException
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode.CONNECTION_CLOSED
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode.INTERNAL_ERROR
|
||||
import kotlin.concurrent.atomics.AtomicReference
|
||||
import kotlin.concurrent.atomics.ExperimentalAtomicApi
|
||||
import kotlin.coroutines.cancellation.CancellationException
|
||||
|
||||
/**
|
||||
* Abstract base class representing a client-side transport layer for communication.
|
||||
*
|
||||
* This class is responsible for managing the lifecycle of a transport instance, including its initialization,
|
||||
* state transitions, message transmission, and graceful shutdown. It provides a framework for implementing
|
||||
* specific transport mechanisms by defining abstract methods that subclasses must implement.
|
||||
*
|
||||
* Thread-safety is achieved through the use of atomic references to manage transport state transitions and
|
||||
* to ensure invariants during lifecycle events. The initialization and shutdown logic is designed to handle
|
||||
* async workflows, ensuring proper resource cleanup in the case of failure or cancellation.
|
||||
*/
|
||||
@OptIn(ExperimentalAtomicApi::class, InternalMcpApi::class)
|
||||
public abstract class AbstractClientTransport : AbstractTransport() {
|
||||
|
||||
protected abstract val logger: KLogger
|
||||
|
||||
/**
|
||||
* Represents the current state of the client transport in an atomic and thread-safe manner.
|
||||
*
|
||||
* This private field is used internally to manage state transitions of the transport, ensuring
|
||||
* consistency across concurrent operations. The state transitions are validated and managed
|
||||
* using the `stateTransition` method, adhering to defined rules in [ClientTransportState.VALID_TRANSITIONS].
|
||||
*
|
||||
* @see ClientTransportState
|
||||
* @see stateTransition
|
||||
* @see checkState
|
||||
*/
|
||||
private val _state: AtomicReference<ClientTransportState> = AtomicReference(ClientTransportState.New)
|
||||
|
||||
/**
|
||||
* Current transport state. Thread-safe via atomic reference.
|
||||
* Initial state is [ClientTransportState.New].
|
||||
*/
|
||||
@InternalMcpApi
|
||||
protected val state: ClientTransportState
|
||||
get() = _state.load()
|
||||
|
||||
/**
|
||||
* Updates the internal state of the client transport to the specified new state.
|
||||
*
|
||||
* This method is used to manage the lifecycle and operational state of the client transport
|
||||
* by transitioning it to the provided [new] state. The new state is stored atomically,
|
||||
* ensuring thread safety during concurrent modifications.
|
||||
*
|
||||
* @param new The new state to transition the client transport to. It must be a valid state
|
||||
* as defined in the transport's lifecycle logic.
|
||||
*/
|
||||
@InternalMcpApi
|
||||
protected fun updateState(new: ClientTransportState) {
|
||||
_state.store(new)
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs transport-specific initialization (connections, I/O streams, background coroutines).
|
||||
*
|
||||
* Called by [start] during [ClientTransportState.Initializing]. On exception, state transitions to
|
||||
* [ClientTransportState.InitializationFailed], [closeResources] is called, and exception propagates.
|
||||
*/
|
||||
protected abstract suspend fun initialize()
|
||||
|
||||
/**
|
||||
* Atomically transitions state from `from` to `to`.
|
||||
*
|
||||
* Validates transition per [ClientTransportState.VALID_TRANSITIONS] and performs atomic compare-and-exchange.
|
||||
*
|
||||
* @throws IllegalArgumentException if transition isn't allowed
|
||||
* @throws IllegalStateException if the actual state doesn't match `from`
|
||||
*/
|
||||
@InternalMcpApi
|
||||
protected fun stateTransition(from: ClientTransportState, to: ClientTransportState) {
|
||||
require(to in ClientTransportState.VALID_TRANSITIONS.getValue(from)) {
|
||||
"Invalid transition: $from → $to"
|
||||
}
|
||||
val actualState = _state.compareAndExchange(
|
||||
expectedValue = from,
|
||||
newValue = to,
|
||||
)
|
||||
check(actualState == from) {
|
||||
"Can't change state: expected transport state $from, but found $actualState."
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Verifies that the current transport state matches the expected state and throws an exception if it does not.
|
||||
*
|
||||
* This method ensures that the transport's state is in alignment with the caller’s expectations before
|
||||
* proceeding with further operations. It is often used to enforce state invariants within the transport's
|
||||
* lifecycle mechanics, enabling robust error handling and debugging in case of unexpected state transitions.
|
||||
*
|
||||
* @param expected The expected transport state. An exception will be thrown if the actual state does not match this value.
|
||||
* @throws IllegalStateException If the current state does not match the expected state.
|
||||
*/
|
||||
protected fun checkState(
|
||||
expected: ClientTransportState,
|
||||
lazyMessage: (ClientTransportState) -> Any = {
|
||||
"Expected transport state $expected, but found $it"
|
||||
},
|
||||
) {
|
||||
val actualState = state
|
||||
check(actualState == expected) { lazyMessage }
|
||||
}
|
||||
|
||||
/**
|
||||
* Initiates the transport by transitioning through the required states to prepare it for operation.
|
||||
*
|
||||
* This method transitions the transport state from `New` to `Initializing`, performs initialization
|
||||
* logic specific to the transport implementation, and then transitions the state to `Operational`
|
||||
* upon successful initialization. If an exception occurs during initialization, the state is transitioned
|
||||
* to `InitializationFailed`, resources are cleaned up, and the exception is rethrown.
|
||||
*
|
||||
* This method operates within a suspend context to allow for asynchronous initialization tasks.
|
||||
*
|
||||
* **Important! This method will eventually become `final`.
|
||||
* Please don't override it, or consistency guarantees might be lost. Override [initialize] instead.**
|
||||
*
|
||||
* @throws Exception if the initialization process fails and transfers state
|
||||
* to [ClientTransportState.InitializationFailed].
|
||||
* @see initialize
|
||||
*/
|
||||
public override suspend fun start() {
|
||||
stateTransition(from = ClientTransportState.New, to = ClientTransportState.Initializing)
|
||||
try {
|
||||
initialize()
|
||||
stateTransition(from = ClientTransportState.Initializing, to = ClientTransportState.Operational)
|
||||
} catch (e: Exception) {
|
||||
_state.store(ClientTransportState.InitializationFailed)
|
||||
closeResources()
|
||||
throw e
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a JSON-RPC message using the transport layer.
|
||||
*
|
||||
* This method ensures that the transport is operational before attempting to send the message.
|
||||
* If the transport's state does not allow sending operations, an exception is thrown.
|
||||
* It delegates the actual transmission to the abstract [performSend] method, which must be
|
||||
* implemented by subclasses to handle the specifics of the transmission mechanism. Errors
|
||||
* during the send operation are handled appropriately, and cancellation exceptions are always
|
||||
* propagated.
|
||||
*
|
||||
* **Important! This method will eventually become `final`.
|
||||
* Please don't override it, or consistency guarantees might be lost. Override [performSend] instead.**
|
||||
*
|
||||
* @param message The JSON-RPC message to be sent. Must adhere to the JSON-RPC 2.0 specification.
|
||||
* @param options Optional parameters that influence how the message is transmitted, such as
|
||||
* timeout specifications or message delivery guarantees. Can be null if no
|
||||
* specific options are required.
|
||||
* @throws McpException If the transport is not operational or other application-specific errors occur.
|
||||
* @throws CancellationException If the coroutine is canceled.
|
||||
* @see performSend
|
||||
*/
|
||||
public override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
// fast path - state check to avoid nonsense operations
|
||||
if (state != ClientTransportState.Operational) {
|
||||
throw McpException(
|
||||
code = CONNECTION_CLOSED,
|
||||
message = "Transport is not ready",
|
||||
)
|
||||
}
|
||||
// limitation: state could still change at the time of use
|
||||
try {
|
||||
performSend(message, options)
|
||||
} catch (e: CancellationException) {
|
||||
throw e // Always propagate cancellation
|
||||
} catch (e: Exception) {
|
||||
_onError(e)
|
||||
if (e is McpException) {
|
||||
throw e
|
||||
} else {
|
||||
throw McpException(
|
||||
code = INTERNAL_ERROR,
|
||||
message = "Error while sending message: ${e.message}",
|
||||
cause = e,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Transmits a JSON-RPC message via the underlying protocol.
|
||||
*
|
||||
* The method is responsible for transmitting a given JSON-RPC message to the intended destination.
|
||||
* It is implemented by concrete subclasses of the transport to handle specific transmission
|
||||
* mechanisms.
|
||||
*
|
||||
* Called by [send] after verifying transport [state].
|
||||
*
|
||||
* @param message The JSON-RPC message to be sent. Must comply with the JSON-RPC 2.0 specification.
|
||||
* @param options Optional parameters that influence how the message is transmitted, such as
|
||||
* timeout settings or delivery guarantees. May be null if no specific options
|
||||
* are required.
|
||||
* @see send
|
||||
*/
|
||||
protected abstract suspend fun performSend(message: JSONRPCMessage, options: TransportSendOptions? = null)
|
||||
|
||||
/**
|
||||
* Closes the client transport gracefully, transitioning its state to a non-operational state.
|
||||
*
|
||||
* **Important! This method will eventually become `final`.
|
||||
* Please don't override it, or consistency guarantees might be lost.**
|
||||
*
|
||||
* The method is designed to operate within a suspendable context, allowing support for asynchronous
|
||||
* behavior during the resource cleanup and state transitions.
|
||||
*
|
||||
* Transitions [ClientTransportState.Operational] → [ClientTransportState.ShuttingDown] → [ClientTransportState.Stopped].
|
||||
* Calls [closeResources] then invokes [_onClose] callbacks exactly once.
|
||||
* On [kotlin.coroutines.cancellation.CancellationException], transitions to [ClientTransportState.ShutdownFailed] and propagates.
|
||||
*
|
||||
* The method also handles unexpected exceptions gracefully during resource cleanup by transitioning
|
||||
* the transport to the `ShutdownFailed` state.
|
||||
*
|
||||
* **Important! This method will eventually become `final`.
|
||||
* Please don't override it, or consistency guarantees might be lost. Override [closeResources] instead.**
|
||||
*
|
||||
* @see closeResources
|
||||
*/
|
||||
public override suspend fun close() {
|
||||
val performClose: Boolean
|
||||
when (state) {
|
||||
ClientTransportState.Operational -> {
|
||||
// Only Operational state can transition to ShuttingDown
|
||||
stateTransition(ClientTransportState.Operational, ClientTransportState.ShuttingDown)
|
||||
performClose = true
|
||||
}
|
||||
|
||||
ClientTransportState.New -> {
|
||||
// New state transitions directly to Stopped without any cleanup
|
||||
stateTransition(ClientTransportState.New, ClientTransportState.Stopped)
|
||||
performClose = false
|
||||
}
|
||||
|
||||
else -> {
|
||||
// Already shutting down or in terminal state - idempotent, do nothing
|
||||
performClose = false
|
||||
}
|
||||
}
|
||||
if (performClose) {
|
||||
try {
|
||||
closeResources()
|
||||
stateTransition(from = ClientTransportState.ShuttingDown, to = ClientTransportState.Stopped)
|
||||
} catch (e: CancellationException) {
|
||||
stateTransition(from = ClientTransportState.ShuttingDown, to = ClientTransportState.ShutdownFailed)
|
||||
throw e // Always propagate cancellation
|
||||
} catch (e: Exception) {
|
||||
// Ignore errors during cleanup
|
||||
logger.error(e) { "Error during transport shutdown" }
|
||||
stateTransition(from = ClientTransportState.ShuttingDown, to = ClientTransportState.ShutdownFailed)
|
||||
} finally {
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Releases transport-specific resources (coroutines, connections, streams).
|
||||
*
|
||||
* Called by [close] during [ClientTransportState.ShuttingDown].
|
||||
* May be called even if [initialize] failed partially.
|
||||
*/
|
||||
protected abstract suspend fun closeResources()
|
||||
}
|
||||
|
|
@ -2,15 +2,18 @@ package io.modelcontextprotocol.kotlin.sdk.shared
|
|||
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlin.concurrent.atomics.AtomicBoolean
|
||||
import kotlin.concurrent.atomics.ExperimentalAtomicApi
|
||||
|
||||
/**
|
||||
* Implements [onClose], [onError] and [onMessage] functions of [Transport] providing
|
||||
* corresponding [_onClose], [_onError] and [_onMessage] properties to use for an implementation.
|
||||
*/
|
||||
@OptIn(ExperimentalAtomicApi::class)
|
||||
@Suppress("PropertyName")
|
||||
public abstract class AbstractTransport : Transport {
|
||||
protected var _onClose: (() -> Unit) = {}
|
||||
private set
|
||||
private val onCloseCalled = AtomicBoolean(false)
|
||||
private var _onClose: (() -> Unit) = {}
|
||||
protected var _onError: ((Throwable) -> Unit) = {}
|
||||
private set
|
||||
|
||||
|
|
@ -51,4 +54,18 @@ public abstract class AbstractTransport : Transport {
|
|||
|
||||
_onMessageInitialized.complete(Unit)
|
||||
}
|
||||
|
||||
/**
|
||||
* Invokes the `_onClose` callback if it has not been already triggered.
|
||||
*
|
||||
* This method ensures the `_onClose` callback is executed only once by utilizing
|
||||
* an atomic flag (`onCloseCalled`). If the callback has already been executed,
|
||||
* the method does nothing. Any exceptions thrown during the execution of the
|
||||
* `_onClose` callback are caught and suppressed.
|
||||
*/
|
||||
protected fun invokeOnCloseCallback() {
|
||||
if (onCloseCalled.compareAndSet(expectedValue = false, newValue = true)) {
|
||||
runCatching { _onClose() }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,79 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
/**
|
||||
* States of a client transport within the protocol lifecycle
|
||||
*
|
||||
* These states are generally aligned with the corresponding phases
|
||||
* defined under the MCP [Protocol Lifecycle](https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle).
|
||||
*/
|
||||
public enum class ClientTransportState {
|
||||
|
||||
/**
|
||||
* The transport has just been created.
|
||||
*
|
||||
* This is the Initial state
|
||||
*/
|
||||
New,
|
||||
|
||||
/**
|
||||
* Indicates that the protocol
|
||||
* is in the [Initialization Phase](https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#initialization).
|
||||
*/
|
||||
Initializing,
|
||||
|
||||
/**
|
||||
* Indicates that the initialization phase of the protocol lifecycle has failed.
|
||||
*
|
||||
* This state suggests that the transport encountered an error or issue during
|
||||
* initialization, preventing it from successfully transitioning to a fully operational state.
|
||||
*/
|
||||
InitializationFailed,
|
||||
|
||||
/**
|
||||
* Represents the [operational phase](https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#operation)
|
||||
* of the protocol lifecycle.
|
||||
*
|
||||
* During the operation phase, the client and server exchange messages according to the negotiated capabilities.
|
||||
*/
|
||||
Operational,
|
||||
|
||||
/**
|
||||
* Represents the shutting down phase of the protocol lifecycle.
|
||||
*
|
||||
* This state indicates that the transport is in the process of shutting down
|
||||
* and transitioning to a non-operational state.
|
||||
*
|
||||
* During this phase, no new outgoing messages should be accepted.
|
||||
*/
|
||||
ShuttingDown,
|
||||
|
||||
/**
|
||||
* Indicates that the shutdown phase of the protocol lifecycle has failed.
|
||||
*
|
||||
* This state signifies that an error was encountered during the shutdown phase.
|
||||
*/
|
||||
ShutdownFailed,
|
||||
|
||||
/**
|
||||
* Indicates that the transport has fully stopped.
|
||||
*
|
||||
* This state represents the final phase of the lifecycle, where the transport is no longer
|
||||
* operational and all cleanup routines or resource deallocations have been completed.
|
||||
*/
|
||||
Stopped,
|
||||
|
||||
;
|
||||
|
||||
internal companion object {
|
||||
val VALID_TRANSITIONS: Map<ClientTransportState, Set<ClientTransportState>> = mapOf(
|
||||
New to setOf(Initializing, Stopped),
|
||||
Initializing to setOf(Operational, InitializationFailed),
|
||||
Operational to setOf(ShuttingDown),
|
||||
ShuttingDown to setOf(Stopped, ShutdownFailed),
|
||||
// Terminal states allow no transitions
|
||||
InitializationFailed to emptySet(),
|
||||
Stopped to emptySet(),
|
||||
ShutdownFailed to emptySet(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -255,7 +255,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
}
|
||||
|
||||
logger.info { "Starting transport" }
|
||||
return transport.start()
|
||||
transport.start()
|
||||
}
|
||||
|
||||
private fun doClose() {
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ public interface Transport {
|
|||
*
|
||||
* This method should only be called after callbacks are installed, or else messages may be lost.
|
||||
*
|
||||
* Method waits for the initialization flow to complete and then returns, so the transport is ready to send and receive user messages.
|
||||
*
|
||||
* NOTE: This method should not be called explicitly when using Client, Server, or Protocol classes,
|
||||
* as they will implicitly call start().
|
||||
*/
|
||||
|
|
@ -26,7 +28,7 @@ public interface Transport {
|
|||
public suspend fun send(message: JSONRPCMessage, options: TransportSendOptions? = null)
|
||||
|
||||
/**
|
||||
* Closes the connection.
|
||||
* Closes the transport.
|
||||
*/
|
||||
public suspend fun close()
|
||||
|
||||
|
|
|
|||
|
|
@ -85,7 +85,7 @@ public abstract class WebSocketMcpTransport : AbstractTransport() {
|
|||
if (it != null) {
|
||||
_onError.invoke(it)
|
||||
} else {
|
||||
_onClose.invoke()
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue