Skip to content

Commit 953acce

Browse files
kpavlovkpavlov
committed
Refactor MCP Protocol and related classes for improved exception handling and annotations.
- Encapsulate exception handling with `toMcpException`. - Move `AbstractTransport` to a separate file. - Add `null` as the default value for optional fields in MCP types. - Rename `McpError` to `McpException` for clarity and better usability. Move to a separate file. Use it in some places instead of IllegalStateException - Gradle: Downgrade `kotest` assertions library to support JDK 8. - Improve `ServerPromptsTest` cases with new scenarios and assertions. - Reformat KDocs and fix Detekt issues - Add additional `@Suppress` annotations to improve readability and control warnings.
1 parent e4bd866 commit 953acce

File tree

13 files changed

+331
-158
lines changed

13 files changed

+331
-158
lines changed

‎build.gradle.kts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ kover {
4646
}
4747
verify {
4848
rule {
49-
minBound(65)
49+
minBound(75)
5050
}
5151
}
5252
}

‎gradle/libs.versions.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ kotlinx-io = "0.8.0"
1818
ktor = "3.2.3"
1919
logging = "7.0.13"
2020
slf4j = "2.0.17"
21-
kotest = "6.0.4"
21+
kotest = "5.9.1" # for JVM 1.8
2222
awaitility = "4.3.0"
2323
mokksy = "0.6.1"
2424

‎kotlin-sdk-core/api/kotlin-sdk-core.api‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
public final class io/modelcontextprotocol/kotlin/sdk/Annotations {
22
public static final field Companion Lio/modelcontextprotocol/kotlin/sdk/Annotations$Companion;
3+
public fun <init> ()V
34
public fun <init> (Ljava/util/List;Lkotlin/time/Instant;Ljava/lang/Double;)V
5+
public synthetic fun <init> (Ljava/util/List;Lkotlin/time/Instant;Ljava/lang/Double;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
46
public final fun component1 ()Ljava/util/List;
57
public final fun component2 ()Lkotlin/time/Instant;
68
public final fun component3 ()Ljava/lang/Double;
@@ -1589,12 +1591,11 @@ public final class io/modelcontextprotocol/kotlin/sdk/LoggingMessageNotification
15891591
public final fun serializer ()Lkotlinx/serialization/KSerializer;
15901592
}
15911593

1592-
public final class io/modelcontextprotocol/kotlin/sdk/McpError : java/lang/Exception {
1594+
public final class io/modelcontextprotocol/kotlin/sdk/McpException : java/lang/Exception {
15931595
public fun <init> (ILjava/lang/String;Lkotlinx/serialization/json/JsonObject;)V
15941596
public synthetic fun <init> (ILjava/lang/String;Lkotlinx/serialization/json/JsonObject;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
15951597
public final fun getCode ()I
15961598
public final fun getData ()Lkotlinx/serialization/json/JsonObject;
1597-
public fun getMessage ()Ljava/lang/String;
15981599
}
15991600

16001601
public abstract interface class io/modelcontextprotocol/kotlin/sdk/Method {
@@ -1699,7 +1700,9 @@ public final class io/modelcontextprotocol/kotlin/sdk/ModelHint$Companion {
16991700

17001701
public final class io/modelcontextprotocol/kotlin/sdk/ModelPreferences {
17011702
public static final field Companion Lio/modelcontextprotocol/kotlin/sdk/ModelPreferences$Companion;
1703+
public fun <init> ()V
17021704
public fun <init> (Ljava/util/List;Ljava/lang/Double;Ljava/lang/Double;Ljava/lang/Double;)V
1705+
public synthetic fun <init> (Ljava/util/List;Ljava/lang/Double;Ljava/lang/Double;Ljava/lang/Double;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
17031706
public final fun getCostPriority ()Ljava/lang/Double;
17041707
public final fun getHints ()Ljava/util/List;
17051708
public final fun getIntelligencePriority ()Ljava/lang/Double;
@@ -1883,6 +1886,7 @@ public final class io/modelcontextprotocol/kotlin/sdk/ProgressNotification$Param
18831886
public final class io/modelcontextprotocol/kotlin/sdk/Prompt {
18841887
public static final field Companion Lio/modelcontextprotocol/kotlin/sdk/Prompt$Companion;
18851888
public fun <init> (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;)V
1889+
public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
18861890
public final fun getArguments ()Ljava/util/List;
18871891
public final fun getDescription ()Ljava/lang/String;
18881892
public final fun getName ()Ljava/lang/String;
@@ -2472,6 +2476,7 @@ public final class io/modelcontextprotocol/kotlin/sdk/Role$Companion {
24722476
public final class io/modelcontextprotocol/kotlin/sdk/Root {
24732477
public static final field Companion Lio/modelcontextprotocol/kotlin/sdk/Root$Companion;
24742478
public fun <init> (Ljava/lang/String;Ljava/lang/String;)V
2479+
public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
24752480
public final fun component1 ()Ljava/lang/String;
24762481
public final fun component2 ()Ljava/lang/String;
24772482
public final fun copy (Ljava/lang/String;Ljava/lang/String;)Lio/modelcontextprotocol/kotlin/sdk/Root;
@@ -2995,6 +3000,7 @@ public final class io/modelcontextprotocol/kotlin/sdk/Tool$Output$Companion {
29953000

29963001
public final class io/modelcontextprotocol/kotlin/sdk/ToolAnnotations {
29973002
public static final field Companion Lio/modelcontextprotocol/kotlin/sdk/ToolAnnotations$Companion;
3003+
public fun <init> ()V
29983004
public fun <init> (Ljava/lang/String;Ljava/lang/Boolean;Ljava/lang/Boolean;Ljava/lang/Boolean;Ljava/lang/Boolean;)V
29993005
public synthetic fun <init> (Ljava/lang/String;Ljava/lang/Boolean;Ljava/lang/Boolean;Ljava/lang/Boolean;Ljava/lang/Boolean;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
30003006
public final fun component1 ()Ljava/lang/String;

‎kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/exceptions.kt‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
@file:Suppress("unused", "EnumEntryName")
2+
3+
package io.modelcontextprotocol.kotlin.sdk
4+
5+
import kotlinx.serialization.json.JsonObject
6+
7+
@Deprecated("Use McpException instead", ReplaceWith("McpException"))
8+
public typealias McpError = McpException
9+
10+
/**
11+
* Represents an error specific to the MCP protocol.
12+
*
13+
* @property code The error code.
14+
* @property message The error message.
15+
* @property data Additional error data as a JSON object.
16+
*/
17+
public class McpException(public val code: Int, message: String, public val data: JsonObject = EmptyJsonObject) :
18+
Exception("MCP error $code: \"$message\"")
19+
20+
/**
21+
* Converts a `JSONRPCError` instance to an [McpException] instance.
22+
*
23+
* @return An [McpException] containing the code, message, and data from the `JSONRPCError`.
24+
*/
25+
internal fun JSONRPCError.toMcpException(): McpException = McpException(
26+
code = this.code.code,
27+
message = this.message,
28+
data = this.data,
29+
)

‎kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/AbstractTransport.kt‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package io.modelcontextprotocol.kotlin.sdk.shared
2+
3+
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
4+
import kotlinx.coroutines.CompletableDeferred
5+
6+
/**
7+
* Implements [onClose], [onError] and [onMessage] functions of [Transport] providing
8+
* corresponding [_onClose], [_onError] and [_onMessage] properties to use for an implementation.
9+
*/
10+
@Suppress("PropertyName")
11+
public abstract class AbstractTransport : Transport {
12+
protected var _onClose: (() -> Unit) = {}
13+
private set
14+
protected var _onError: ((Throwable) -> Unit) = {}
15+
private set
16+
17+
// to not skip messages
18+
private val _onMessageInitialized = CompletableDeferred<Unit>()
19+
protected var _onMessage: (suspend ((JSONRPCMessage) -> Unit)) = {
20+
_onMessageInitialized.await()
21+
_onMessage.invoke(it)
22+
}
23+
private set
24+
25+
override fun onClose(block: () -> Unit) {
26+
val old = _onClose
27+
_onClose = {
28+
old()
29+
block()
30+
}
31+
}
32+
33+
override fun onError(block: (Throwable) -> Unit) {
34+
val old = _onError
35+
_onError = { e ->
36+
old(e)
37+
block(e)
38+
}
39+
}
40+
41+
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {
42+
val old: suspend (JSONRPCMessage) -> Unit = when (_onMessageInitialized.isCompleted) {
43+
true -> _onMessage
44+
false -> { _ -> }
45+
}
46+
47+
_onMessage = { message ->
48+
old(message)
49+
block(message)
50+
}
51+
52+
_onMessageInitialized.complete(Unit)
53+
}
54+
}

‎kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt‎

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import io.modelcontextprotocol.kotlin.sdk.JSONRPCError
88
import io.modelcontextprotocol.kotlin.sdk.JSONRPCNotification
99
import io.modelcontextprotocol.kotlin.sdk.JSONRPCRequest
1010
import io.modelcontextprotocol.kotlin.sdk.JSONRPCResponse
11-
import io.modelcontextprotocol.kotlin.sdk.McpError
11+
import io.modelcontextprotocol.kotlin.sdk.McpException
1212
import io.modelcontextprotocol.kotlin.sdk.Method
1313
import io.modelcontextprotocol.kotlin.sdk.Notification
1414
import io.modelcontextprotocol.kotlin.sdk.PingRequest
@@ -19,6 +19,7 @@ import io.modelcontextprotocol.kotlin.sdk.RequestId
1919
import io.modelcontextprotocol.kotlin.sdk.RequestResult
2020
import io.modelcontextprotocol.kotlin.sdk.fromJSON
2121
import io.modelcontextprotocol.kotlin.sdk.toJSON
22+
import io.modelcontextprotocol.kotlin.sdk.toMcpException
2223
import kotlinx.atomicfu.AtomicRef
2324
import kotlinx.atomicfu.atomic
2425
import kotlinx.atomicfu.getAndUpdate
@@ -97,7 +98,7 @@ public data class RequestOptions(
9798
val onProgress: ProgressCallback? = null,
9899

99100
/**
100-
* A timeout for this request. If exceeded, an McpError with code `RequestTimeout`
101+
* A timeout for this request. If exceeded, an [McpException] with code `RequestTimeout`
101102
* will be raised from request().
102103
*
103104
* If not specified, `DEFAULT_REQUEST_TIMEOUT` will be used as the timeout.
@@ -116,6 +117,7 @@ internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) }
116117
* Implements MCP protocol framing on top of a pluggable transport, including
117118
* features like request/response linking, notifications, and progress.
118119
*/
120+
@Suppress("TooManyFunctions")
119121
public abstract class Protocol(@PublishedApi internal val options: ProtocolOptions?) {
120122
public var transport: Transport? = null
121123
private set
@@ -190,7 +192,9 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
190192
/**
191193
* Attaches to the given transport, starts it, and starts listening for messages.
192194
*
193-
* The Protocol object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward.
195+
* The Protocol object assumes ownership of the Transport,
196+
* replacing any callbacks that have already been set,
197+
* and expects that it is the only user of the Transport instance going forward.
194198
*/
195199
public open suspend fun connect(transport: Transport) {
196200
this.transport = transport
@@ -222,7 +226,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
222226
transport = null
223227
onClose()
224228

225-
val error = McpError(ErrorCode.Defined.ConnectionClosed.code, "Connection closed")
229+
val error = McpException(ErrorCode.Defined.ConnectionClosed.code, "Connection closed")
226230
for (handler in handlersToNotify) {
227231
handler(null, error)
228232
}
@@ -237,6 +241,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
237241
logger.trace { "No handler found for notification: ${notification.method}" }
238242
return
239243
}
244+
@Suppress("TooGenericExceptionCaught")
240245
try {
241246
handler(notification)
242247
} catch (cause: Throwable) {
@@ -252,6 +257,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
252257

253258
if (handler === null) {
254259
logger.trace { "No handler found for request: ${request.method}" }
260+
@Suppress("TooGenericExceptionCaught")
255261
try {
256262
transport?.send(
257263
JSONRPCResponse(
@@ -268,7 +274,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
268274
}
269275
return
270276
}
271-
277+
@Suppress("TooGenericExceptionCaught")
272278
try {
273279
val result = handler(request, RequestHandlerExtra())
274280
logger.trace { "Request handled successfully: ${request.method} (id: ${request.id})" }
@@ -303,7 +309,8 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
303309

304310
private fun onProgress(notification: ProgressNotification) {
305311
logger.trace {
306-
"Received progress notification: token=${notification.params.progressToken}, progress=${notification.params.progress}/${notification.params.total}"
312+
"Received progress notification: token=${notification.params.progressToken}, " +
313+
"progress=${notification.params.progress}/${notification.params.total}"
307314
}
308315
val progress = notification.params.progress
309316
val total = notification.params.total
@@ -347,7 +354,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
347354
handler(response, null)
348355
} else {
349356
check(error != null)
350-
val error = McpError(
357+
val error = McpException(
351358
error.code.code,
352359
error.message,
353360
error.data,
@@ -392,7 +399,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
392399
public suspend fun <T : RequestResult> request(request: Request, options: RequestOptions? = null): T {
393400
logger.trace { "Sending request: ${request.method}" }
394401
val result = CompletableDeferred<T>()
395-
val transport = transport ?: throw Error("Not connected")
402+
val transport = transport ?: throw McpException(ErrorCode.Defined.ConnectionClosed.code, "Not connected")
396403

397404
if (this@Protocol.options?.enforceStrictCapabilities == true) {
398405
assertCapabilityForMethod(request.method)
@@ -415,11 +422,12 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
415422
return@put
416423
}
417424

418-
if (response?.error != null) {
419-
result.completeExceptionally(IllegalStateException(response.error.toString()))
425+
response?.error?.let {
426+
result.completeExceptionally(it.toMcpException())
420427
return@put
421428
}
422429

430+
@Suppress("TooGenericExceptionCaught")
423431
try {
424432
@Suppress("UNCHECKED_CAST")
425433
result.complete(response!!.result as T)
@@ -459,7 +467,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
459467
} catch (cause: TimeoutCancellationException) {
460468
logger.error { "Request timed out after ${timeout.inWholeMilliseconds}ms: ${request.method}" }
461469
cancel(
462-
McpError(
470+
McpException(
463471
ErrorCode.Defined.RequestTimeout.code,
464472
"Request timed out",
465473
JsonObject(mutableMapOf("timeout" to JsonPrimitive(timeout.inWholeMilliseconds))),

‎kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Transport.kt‎

Lines changed: 0 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package io.modelcontextprotocol.kotlin.sdk.shared
22

33
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
4-
import kotlinx.coroutines.CompletableDeferred
54

65
/**
76
* Describes the minimal contract for MCP transport that a client or server can communicate over.
@@ -47,53 +46,3 @@ public interface Transport {
4746
*/
4847
public fun onMessage(block: suspend (JSONRPCMessage) -> Unit)
4948
}
50-
51-
/**
52-
* Implements [onClose], [onError] and [onMessage] functions of [Transport] providing
53-
* corresponding [_onClose], [_onError] and [_onMessage] properties to use for an implementation.
54-
*/
55-
@Suppress("PropertyName")
56-
public abstract class AbstractTransport : Transport {
57-
protected var _onClose: (() -> Unit) = {}
58-
private set
59-
protected var _onError: ((Throwable) -> Unit) = {}
60-
private set
61-
62-
// to not skip messages
63-
private val _onMessageInitialized = CompletableDeferred<Unit>()
64-
protected var _onMessage: (suspend ((JSONRPCMessage) -> Unit)) = {
65-
_onMessageInitialized.await()
66-
_onMessage.invoke(it)
67-
}
68-
private set
69-
70-
override fun onClose(block: () -> Unit) {
71-
val old = _onClose
72-
_onClose = {
73-
old()
74-
block()
75-
}
76-
}
77-
78-
override fun onError(block: (Throwable) -> Unit) {
79-
val old = _onError
80-
_onError = { e ->
81-
old(e)
82-
block(e)
83-
}
84-
}
85-
86-
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {
87-
val old: suspend (JSONRPCMessage) -> Unit = when (_onMessageInitialized.isCompleted) {
88-
true -> _onMessage
89-
false -> { _ -> }
90-
}
91-
92-
_onMessage = { message ->
93-
old(message)
94-
block(message)
95-
}
96-
97-
_onMessageInitialized.complete(Unit)
98-
}
99-
}

0 commit comments

Comments
 (0)