feat(core): concurrent inbound dispatch behind an initialization gate (#884)
Make inbound message handling concurrent after the initialization handshake, so a slow or peer-awaiting handler no longer blocks other messages on the same connection fixes #176, #572 ## Motivation and Context Inbound messages were processed serially on the transport read loop. One slow or peer-awaiting request handler (sampling, elicitation, roots) blocked every later message on that connection, including the responses and `notifications/cancelled` those handlers depend on which could deadlock request/response flows. ## How Has This Been Tested? New and old tests ## Breaking Changes - after the handshake, inbound handlers may run concurrently (processing stays serial during initialization). There is no opt-out flag by design, concurrency is bounded by `maxConcurrentHandlers` / `maxInFlightHandlers`. - `ProtocolOptions` / `ClientOptions` / `ServerOptions` gain new parameters with defaults - `RequestHandlerExtra` is now a coroutine-context element and is no longer user-constructible ## Types of changes - [x] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update ## Checklist - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [x] New and existing tests pass locally - [x] I have added appropriate error handling - [x] I have added or updated documentation as needed
This commit is contained in:
Pavel Gorgulov
GitHub
parent
7937390b47
commit
beb3400064
21 changed files with 2462 additions and 191 deletions
|
|
@ -24,6 +24,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest
|
|||
import io.modelcontextprotocol.kotlin.sdk.types.InitializeResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.IntegerSchema
|
||||
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.LATEST_PROTOCOL_VERSION
|
||||
|
|
@ -52,6 +53,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.UntitledSingleSelectEnumSchema
|
|||
import io.modelcontextprotocol.kotlin.sdk.types.UrlElicitationRequiredException
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.TimeoutCancellationException
|
||||
import kotlinx.coroutines.awaitCancellation
|
||||
import kotlinx.coroutines.cancel
|
||||
import kotlinx.coroutines.delay
|
||||
import kotlinx.coroutines.joinAll
|
||||
|
|
@ -71,6 +73,7 @@ import kotlin.test.assertIs
|
|||
import kotlin.test.assertNull
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.test.fail
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
class ClientTest {
|
||||
@Test
|
||||
|
|
@ -543,6 +546,15 @@ class ClientTest {
|
|||
|
||||
val serverSession = serverSessionResult.await()
|
||||
|
||||
// Register the tap after createSession so it chains onto the session's subscription instead
|
||||
// of pre-empting the handshake.
|
||||
val cancelledOnWire = CompletableDeferred<Unit>()
|
||||
serverTransport.onMessage { message ->
|
||||
if (message is JSONRPCNotification && message.method == Method.Defined.NotificationsCancelled.value) {
|
||||
cancelledOnWire.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
serverSession.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
// Simulate delay
|
||||
def.complete(Unit)
|
||||
|
|
@ -568,6 +580,7 @@ class ClientTest {
|
|||
def.await()
|
||||
runCatching { job.cancel("Cancelled by test") }
|
||||
defCancel.await()
|
||||
cancelledOnWire.await()
|
||||
defTimeOut.await()
|
||||
}
|
||||
|
||||
|
|
@ -606,17 +619,9 @@ class ClientTest {
|
|||
|
||||
val serverSession = serverSessionResult.await()
|
||||
serverSession.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
// Simulate a delayed response
|
||||
// Wait ~100ms unless canceled
|
||||
try {
|
||||
withTimeout(100L) {
|
||||
// Just delay here, if timeout is 0 on the client side, this won't return in time
|
||||
delay(100)
|
||||
}
|
||||
} catch (_: Exception) {
|
||||
// If aborted, just rethrow or return early
|
||||
}
|
||||
ListResourcesResult(resources = emptyList())
|
||||
// Never respond: the handler runs on a real dispatcher while the 1 ms timeout runs on the
|
||||
// virtual test clock, so a real-time delay could race it. Parking keeps the test deterministic.
|
||||
awaitCancellation()
|
||||
}
|
||||
|
||||
// Request with 1 msec timeout should fail immediately
|
||||
|
|
@ -895,8 +900,8 @@ class ClientTest {
|
|||
),
|
||||
)
|
||||
|
||||
// Track notifications
|
||||
var rootListChangedNotificationReceived = false
|
||||
// The handler runs on Dispatchers.Default after concurrent dispatch is enabled; await the deferred.
|
||||
val rootListChangedReceived = CompletableDeferred<Unit>()
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
|
|
@ -915,16 +920,13 @@ class ClientTest {
|
|||
serverSession.setNotificationHandler<RootsListChangedNotification>(
|
||||
Method.Defined.NotificationsRootsListChanged,
|
||||
) {
|
||||
rootListChangedNotificationReceived = true
|
||||
rootListChangedReceived.complete(Unit)
|
||||
CompletableDeferred(Unit)
|
||||
}
|
||||
|
||||
client.sendRootsListChanged()
|
||||
|
||||
assertTrue(
|
||||
rootListChangedNotificationReceived,
|
||||
"Notification should be sent when sendRootsListChanged is called",
|
||||
)
|
||||
withTimeout(5.seconds) { rootListChangedReceived.await() }
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
|
|||
|
|
@ -0,0 +1,306 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.integration
|
||||
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.modelcontextprotocol.kotlin.sdk.ExperimentalMcpApi
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.Client
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.ClientOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.Server
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.ServerSession
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.InMemoryTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.currentRequestHandlerExtra
|
||||
import io.modelcontextprotocol.kotlin.sdk.testing.ChannelTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CallToolRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CallToolRequestParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageRequestParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequestParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ListRootsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ListRootsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ListToolsRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ListToolsResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Role
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.SamplingMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.TextContent
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.toJSON
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import java.util.concurrent.ConcurrentHashMap
|
||||
import java.util.concurrent.CopyOnWriteArrayList
|
||||
import java.util.concurrent.atomic.AtomicReference
|
||||
import kotlin.coroutines.cancellation.CancellationException
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertFalse
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
/**
|
||||
* End-to-end tests for concurrent inbound dispatch enabled after the MCP handshake.
|
||||
*
|
||||
* Handlers run on a real [kotlinx.coroutines.Dispatchers.Default], so these use [runBlocking] with
|
||||
* explicit [withTimeout] guards. A virtual test clock would false-fire timeouts against cross-thread
|
||||
* real dispatchers.
|
||||
*/
|
||||
@OptIn(ExperimentalMcpApi::class)
|
||||
class ProtocolConcurrencyIntegrationTest {
|
||||
|
||||
private fun rootsClient(): Client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(roots = ClientCapabilities.Roots(listChanged = null)),
|
||||
),
|
||||
)
|
||||
|
||||
private fun toolsServer(): Server = Server(
|
||||
serverInfo = Implementation(name = "test server", version = "1.0"),
|
||||
options = ServerOptions(
|
||||
capabilities = ServerCapabilities(tools = ServerCapabilities.Tools(listChanged = null)),
|
||||
),
|
||||
)
|
||||
|
||||
// ChannelTransport delivers inbound messages as soon as its read loop starts, which is before
|
||||
// start() marks it operational. Creating the session first keeps the client's initialize from
|
||||
// landing on a transport that would reject the session's response.
|
||||
private suspend fun connect(
|
||||
client: Client,
|
||||
clientTransport: ChannelTransport,
|
||||
server: Server,
|
||||
serverTransport: ChannelTransport,
|
||||
): ServerSession {
|
||||
val serverSession = server.createSession(serverTransport)
|
||||
client.connect(clientTransport)
|
||||
return serverSession
|
||||
}
|
||||
|
||||
// The motivating deadlock over real serial read loops: a server tool handler makes a
|
||||
// server-to-client request (roots/list) while still handling the client's tools/call.
|
||||
@Test
|
||||
fun `server tool handler can call roots list mid-handling without deadlock`() = runBlocking {
|
||||
val client = rootsClient()
|
||||
val server = toolsServer()
|
||||
val (clientTransport, serverTransport) = ChannelTransport.createLinkedPair()
|
||||
|
||||
val serverSession = connect(client, clientTransport, server, serverTransport)
|
||||
|
||||
// Register the tools/call handler directly on the session so it can call back to the client
|
||||
// mid-handling. Before concurrent dispatch this would deadlock: the server read loop is
|
||||
// blocked by the tool handler and can never deliver the roots/list response.
|
||||
serverSession.setRequestHandler<CallToolRequest>(Method.Defined.ToolsCall) { _, _ ->
|
||||
val roots = serverSession.request<ListRootsResult>(ListRootsRequest())
|
||||
CallToolResult(content = listOf(TextContent(text = "roots=${roots.roots.size}")))
|
||||
}
|
||||
|
||||
val result = withTimeout(10.seconds) {
|
||||
client.callTool(CallToolRequest(CallToolRequestParams("needs-roots")))
|
||||
}
|
||||
val text = (result as CallToolResult).content.filterIsInstance<TextContent>().single().text
|
||||
text shouldBe "roots=0"
|
||||
|
||||
client.close()
|
||||
server.close()
|
||||
}
|
||||
|
||||
// Client-side mirror: a client sampling handler makes a client-to-server request (tools/list)
|
||||
// while still handling the server's sampling/createMessage.
|
||||
@Test
|
||||
fun `client sampling handler can call the server mid-handling without deadlock`() = runBlocking {
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(sampling = ClientCapabilities.Sampling()),
|
||||
),
|
||||
)
|
||||
val server = toolsServer()
|
||||
val (clientTransport, serverTransport) = ChannelTransport.createLinkedPair()
|
||||
|
||||
client.setRequestHandler<CreateMessageRequest>(Method.Defined.SamplingCreateMessage) { _, _ ->
|
||||
val tools = client.request<ListToolsResult>(ListToolsRequest())
|
||||
CreateMessageResult(
|
||||
role = Role.Assistant,
|
||||
content = TextContent(text = "tools=${tools.tools.size}"),
|
||||
model = "test-model",
|
||||
)
|
||||
}
|
||||
|
||||
val serverSession = connect(client, clientTransport, server, serverTransport)
|
||||
|
||||
val result = withTimeout(10.seconds) {
|
||||
serverSession.request<CreateMessageResult>(
|
||||
CreateMessageRequest(
|
||||
CreateMessageRequestParams(
|
||||
maxTokens = 100,
|
||||
messages = listOf(SamplingMessage(role = Role.User, content = TextContent(text = "hi"))),
|
||||
),
|
||||
),
|
||||
)
|
||||
}
|
||||
val text = result.content.filterIsInstance<TextContent>().single().text
|
||||
text shouldBe "tools=0"
|
||||
|
||||
client.close()
|
||||
server.close()
|
||||
}
|
||||
|
||||
// Nested-cancel cascade: cancelling the client's tools/call must reach the server tool handler
|
||||
// over the wire and cascade into its nested roots/list request, cancelling the parked
|
||||
// client-side roots handler. No tools/call response is ever sent.
|
||||
@Test
|
||||
fun `cancelling a tool call cascades into its nested roots list request`() = runBlocking {
|
||||
val client = rootsClient()
|
||||
val server = toolsServer()
|
||||
val (clientTransport, serverTransport) = ChannelTransport.createLinkedPair()
|
||||
|
||||
val capturedToolCallId = AtomicReference<RequestId?>(null)
|
||||
val toolHandlerCancelled = CompletableDeferred<CancellationException>()
|
||||
val rootsHandlerEntered = CompletableDeferred<Unit>()
|
||||
val rootsHandlerCancelled = CompletableDeferred<CancellationException>()
|
||||
val parkForever = CompletableDeferred<ListRootsResult>()
|
||||
val unexpectedToolCallResponse = CompletableDeferred<Unit>()
|
||||
|
||||
// Tap the server->client direction (client transport inbound) to prove no tools/call
|
||||
// response is ever delivered once the call is cancelled.
|
||||
clientTransport.onMessage { message ->
|
||||
if (message is JSONRPCResponse && message.id == capturedToolCallId.get()) {
|
||||
unexpectedToolCallResponse.complete(Unit)
|
||||
}
|
||||
}
|
||||
|
||||
val serverSession = connect(client, clientTransport, server, serverTransport)
|
||||
|
||||
// Client roots handler parks until cancelled and records the cancellation.
|
||||
client.setRequestHandler<ListRootsRequest>(Method.Defined.RootsList) { _, _ ->
|
||||
rootsHandlerEntered.complete(Unit)
|
||||
try {
|
||||
parkForever.await()
|
||||
} catch (e: CancellationException) {
|
||||
rootsHandlerCancelled.complete(e)
|
||||
throw e
|
||||
}
|
||||
}
|
||||
|
||||
// Tool handler makes a nested server->client request and records its own cancellation.
|
||||
serverSession.setRequestHandler<CallToolRequest>(Method.Defined.ToolsCall) { _, _ ->
|
||||
val extra = currentRequestHandlerExtra()!!
|
||||
capturedToolCallId.set(extra.requestId)
|
||||
try {
|
||||
extra.sendRequest<ListRootsResult>(ListRootsRequest())
|
||||
} catch (e: CancellationException) {
|
||||
toolHandlerCancelled.complete(e)
|
||||
throw e
|
||||
}
|
||||
CallToolResult(content = listOf(TextContent(text = "unreachable")))
|
||||
}
|
||||
|
||||
val call = launch { client.callTool(CallToolRequest(CallToolRequestParams("nested"))) }
|
||||
|
||||
withTimeout(10.seconds) { rootsHandlerEntered.await() }
|
||||
call.cancel(CancellationException("user aborted"))
|
||||
|
||||
withTimeout(10.seconds) { toolHandlerCancelled.await() }
|
||||
withTimeout(10.seconds) { rootsHandlerCancelled.await() }
|
||||
|
||||
// The server suppresses the response for a cancelled request, so none is ever sent.
|
||||
assertFalse(
|
||||
unexpectedToolCallResponse.isCompleted,
|
||||
"no tools/call response should be delivered after the call was cancelled",
|
||||
)
|
||||
|
||||
client.close()
|
||||
server.close()
|
||||
}
|
||||
|
||||
// Init gate plus context element in a high-level handler: initialize and tools/list are
|
||||
// pipelined before notifications/initialized and answered serially in order. After the
|
||||
// handshake a high-level addTool handler can observe currentRequestHandlerExtra().
|
||||
@Test
|
||||
fun `pipelined initialize and tools list are processed serially then context element is available`(): Unit =
|
||||
runBlocking {
|
||||
val server = toolsServer()
|
||||
server.addTool(name = "ctx-tool", description = "returns its request id") {
|
||||
val extra = currentRequestHandlerExtra()
|
||||
checkNotNull(extra) { "currentRequestHandlerExtra() must be non-null inside a tool handler" }
|
||||
CallToolResult(content = listOf(TextContent(text = extra.requestId.asText())))
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val arrivalOrder = CopyOnWriteArrayList<RequestId>()
|
||||
val pending = ConcurrentHashMap<RequestId, CompletableDeferred<JSONRPCMessage>>()
|
||||
clientTransport.onMessage { message ->
|
||||
val id = when (message) {
|
||||
is JSONRPCResponse -> message.id
|
||||
is JSONRPCError -> message.id
|
||||
else -> null
|
||||
}
|
||||
if (id != null) {
|
||||
arrivalOrder.add(id)
|
||||
pending[id]?.complete(message)
|
||||
}
|
||||
}
|
||||
|
||||
// Drive the client side manually over the raw transport.
|
||||
server.createSession(serverTransport)
|
||||
|
||||
val initRequest = InitializeRequest(
|
||||
InitializeRequestParams(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ClientCapabilities(),
|
||||
clientInfo = Implementation(name = "raw-client", version = "1.0"),
|
||||
),
|
||||
).toJSON()
|
||||
val toolsListRequest = ListToolsRequest().toJSON()
|
||||
|
||||
val initAnswered = pending.expect(initRequest.id)
|
||||
val toolsListAnswered = pending.expect(toolsListRequest.id)
|
||||
|
||||
// Pipelined before notifications/initialized: both handled in the serial dispatch phase.
|
||||
clientTransport.send(initRequest)
|
||||
clientTransport.send(toolsListRequest)
|
||||
|
||||
withTimeout(10.seconds) { initAnswered.await() }
|
||||
withTimeout(10.seconds) { toolsListAnswered.await() }
|
||||
|
||||
// Deterministic serial ordering: initialize answered before tools/list.
|
||||
assertEquals(
|
||||
listOf(initRequest.id, toolsListRequest.id),
|
||||
arrivalOrder.toList(),
|
||||
"pipelined requests must be answered in arrival order",
|
||||
)
|
||||
|
||||
// Complete the handshake, then call the high-level tool.
|
||||
clientTransport.send(InitializedNotification().toJSON())
|
||||
|
||||
val toolCallRequest = CallToolRequest(CallToolRequestParams("ctx-tool")).toJSON()
|
||||
val toolCallAnswered = pending.expect(toolCallRequest.id)
|
||||
clientTransport.send(toolCallRequest)
|
||||
|
||||
val toolCallResponse = withTimeout(10.seconds) { toolCallAnswered.await() } as JSONRPCResponse
|
||||
val text = (toolCallResponse.result as CallToolResult).content.filterIsInstance<TextContent>().single().text
|
||||
// The handler echoed its request id, proving currentRequestHandlerExtra() was available.
|
||||
text shouldBe toolCallRequest.id.asText()
|
||||
}
|
||||
|
||||
private fun ConcurrentHashMap<RequestId, CompletableDeferred<JSONRPCMessage>>.expect(
|
||||
id: RequestId,
|
||||
): CompletableDeferred<JSONRPCMessage> = CompletableDeferred<JSONRPCMessage>().also { put(id, it) }
|
||||
|
||||
private fun RequestId.asText(): String = when (this) {
|
||||
is RequestId.StringId -> value
|
||||
is RequestId.NumberId -> value.toString()
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,7 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.kotest.assertions.nondeterministic.continually
|
||||
import io.kotest.assertions.nondeterministic.eventually
|
||||
import io.kotest.matchers.collections.shouldBeEmpty
|
||||
import io.kotest.matchers.collections.shouldHaveSize
|
||||
import io.kotest.matchers.shouldBe
|
||||
|
|
@ -17,6 +19,9 @@ import kotlinx.serialization.json.JsonPrimitive
|
|||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.params.ParameterizedTest
|
||||
import org.junit.jupiter.params.provider.EnumSource
|
||||
import java.util.concurrent.CopyOnWriteArrayList
|
||||
import kotlin.time.Duration.Companion.milliseconds
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
class ClientConnectionLoggingTest : AbstractServerFeaturesTest() {
|
||||
|
||||
|
|
@ -56,7 +61,7 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() {
|
|||
@ParameterizedTest
|
||||
@EnumSource(LoggingLevel::class)
|
||||
fun `notification should filter logging messages below level`(minLevel: LoggingLevel): Unit = runBlocking {
|
||||
val receivedMessages = mutableListOf<LoggingMessageNotification>()
|
||||
val receivedMessages = CopyOnWriteArrayList<LoggingMessageNotification>()
|
||||
client.setNotificationHandler<LoggingMessageNotification>(Method.Defined.NotificationsMessage) {
|
||||
receivedMessages.add(it)
|
||||
CompletableDeferred(Unit)
|
||||
|
|
@ -79,8 +84,10 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() {
|
|||
|
||||
client.callTool(CallToolRequest(CallToolRequestParams("test-notification-filtered")))
|
||||
|
||||
// Handlers dispatch on Dispatchers.Default after the handshake, so await the expected count
|
||||
// before asserting order.
|
||||
val expectedLevels = LoggingLevel.entries.filter { it >= minLevel }
|
||||
receivedMessages shouldHaveSize expectedLevels.size
|
||||
eventually(5.seconds) { receivedMessages shouldHaveSize expectedLevels.size }
|
||||
receivedMessages.map { it.params.level } shouldBe expectedLevels
|
||||
}
|
||||
|
||||
|
|
@ -115,7 +122,7 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() {
|
|||
@ParameterizedTest
|
||||
@EnumSource(LoggingLevel::class)
|
||||
fun `sendLoggingMessage should filter messages below level`(minLevel: LoggingLevel): Unit = runBlocking {
|
||||
val receivedMessages = mutableListOf<LoggingMessageNotification>()
|
||||
val receivedMessages = CopyOnWriteArrayList<LoggingMessageNotification>()
|
||||
client.setNotificationHandler<LoggingMessageNotification>(Method.Defined.NotificationsMessage) {
|
||||
receivedMessages.add(it)
|
||||
CompletableDeferred(Unit)
|
||||
|
|
@ -138,14 +145,16 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() {
|
|||
|
||||
client.callTool(CallToolRequest(CallToolRequestParams("test-logging-level")))
|
||||
|
||||
// Handlers dispatch on Dispatchers.Default after the handshake, so await the expected count
|
||||
// before asserting order.
|
||||
val expectedLevels = LoggingLevel.entries.filter { it >= minLevel }
|
||||
receivedMessages shouldHaveSize expectedLevels.size
|
||||
eventually(5.seconds) { receivedMessages shouldHaveSize expectedLevels.size }
|
||||
receivedMessages.map { it.params.level } shouldBe expectedLevels
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `sendLoggingMessage should send no messages when level is set to highest`(): Unit = runBlocking {
|
||||
val receivedMessages = mutableListOf<LoggingMessageNotification>()
|
||||
val receivedMessages = CopyOnWriteArrayList<LoggingMessageNotification>()
|
||||
client.setNotificationHandler<LoggingMessageNotification>(Method.Defined.NotificationsMessage) {
|
||||
receivedMessages.add(it)
|
||||
CompletableDeferred(Unit)
|
||||
|
|
@ -167,6 +176,7 @@ class ClientConnectionLoggingTest : AbstractServerFeaturesTest() {
|
|||
}
|
||||
|
||||
client.callTool(CallToolRequest(CallToolRequestParams("test-logging-highest")))
|
||||
receivedMessages.shouldBeEmpty()
|
||||
// Handlers dispatch asynchronously, so hold for a window to confirm nothing ever arrives.
|
||||
continually(500.milliseconds) { receivedMessages.shouldBeEmpty() }
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -65,6 +65,8 @@ class ServerSessionInitializeTest {
|
|||
assertEquals("test-client", session.clientVersion?.name)
|
||||
}
|
||||
|
||||
// Both initialize requests arrive before notifications/initialized, in the serial dispatch
|
||||
// phase, so processing and response order stay deterministic.
|
||||
@Test
|
||||
fun `should reject duplicate initialize request`() = runTest {
|
||||
val session = createSession()
|
||||
|
|
|
|||
|
|
@ -52,8 +52,8 @@ public final class io/modelcontextprotocol/kotlin/sdk/client/ClientKt {
|
|||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/client/ClientOptions : io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions {
|
||||
public fun <init> ()V
|
||||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;Z)V
|
||||
public synthetic fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;ZLkotlin/coroutines/CoroutineContext;)V
|
||||
public synthetic fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;ZLkotlin/coroutines/CoroutineContext;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public final fun getCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -73,11 +73,13 @@ import kotlinx.collections.immutable.minus
|
|||
import kotlinx.collections.immutable.persistentMapOf
|
||||
import kotlinx.collections.immutable.toPersistentSet
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.serialization.SerializationException
|
||||
import kotlinx.serialization.json.JsonElement
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonPrimitive
|
||||
import kotlinx.serialization.json.buildJsonArray
|
||||
import kotlin.coroutines.CoroutineContext
|
||||
import kotlin.coroutines.cancellation.CancellationException
|
||||
|
||||
private val logger = KotlinLogging.logger {}
|
||||
|
|
@ -87,11 +89,16 @@ private val logger = KotlinLogging.logger {}
|
|||
*
|
||||
* @property capabilities The capabilities this client supports.
|
||||
* @param enforceStrictCapabilities Whether to strictly enforce capabilities when interacting with the server.
|
||||
* @param handlerCoroutineContext Coroutine context for inbound handlers. See [ProtocolOptions.handlerCoroutineContext].
|
||||
*/
|
||||
public class ClientOptions(
|
||||
public val capabilities: ClientCapabilities = ClientCapabilities(),
|
||||
enforceStrictCapabilities: Boolean = true,
|
||||
) : ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities)
|
||||
handlerCoroutineContext: CoroutineContext = Dispatchers.Default,
|
||||
) : ProtocolOptions(
|
||||
enforceStrictCapabilities = enforceStrictCapabilities,
|
||||
handlerCoroutineContext = handlerCoroutineContext,
|
||||
)
|
||||
|
||||
/**
|
||||
* Initializes and connects an MCP client using the provided clientInfo [Implementation], client options,
|
||||
|
|
@ -216,6 +223,7 @@ public open class Client(private val clientInfo: Implementation, options: Client
|
|||
serverInstructions = result.instructions
|
||||
|
||||
notification(InitializedNotification())
|
||||
enableConcurrentDispatch()
|
||||
} catch (error: Throwable) {
|
||||
logger.error(error) { "Failed to initialize client: ${error.message}" }
|
||||
close()
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol {
|
|||
public abstract fun assertRequestHandlerCapability (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V
|
||||
public final fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public fun connect (Lio/modelcontextprotocol/kotlin/sdk/shared/Transport;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
protected final fun enableConcurrentDispatch ()V
|
||||
public final fun getFallbackNotificationHandler ()Lkotlin/jvm/functions/Function2;
|
||||
public final fun getFallbackRequestHandler ()Lkotlin/jvm/functions/Function3;
|
||||
public final fun getNotificationHandlers ()Ljava/util/Map;
|
||||
|
|
@ -71,6 +72,7 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol {
|
|||
public static synthetic fun notification$default (Lio/modelcontextprotocol/kotlin/sdk/shared/Protocol;Lio/modelcontextprotocol/kotlin/sdk/types/Notification;Lio/modelcontextprotocol/kotlin/sdk/types/RequestId;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
|
||||
public fun onClose ()V
|
||||
public fun onError (Ljava/lang/Throwable;)V
|
||||
protected fun onInitializedNotification ()V
|
||||
public final fun removeNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V
|
||||
public final fun removeRequestHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V
|
||||
public final fun request (Lio/modelcontextprotocol/kotlin/sdk/types/Request;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
|
|
@ -84,15 +86,18 @@ public abstract class io/modelcontextprotocol/kotlin/sdk/shared/Protocol {
|
|||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/shared/ProtocolKt {
|
||||
public static final field IMPLEMENTATION_NAME Ljava/lang/String;
|
||||
public static final fun currentRequestHandlerExtra (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public static final fun getDEFAULT_REQUEST_TIMEOUT ()J
|
||||
}
|
||||
|
||||
public class io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions {
|
||||
public synthetic fun <init> (ZJILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public synthetic fun <init> (ZJLkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public synthetic fun <init> (ZJLkotlin/coroutines/CoroutineContext;IIILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public synthetic fun <init> (ZJLkotlin/coroutines/CoroutineContext;IILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public final fun getEnforceStrictCapabilities ()Z
|
||||
public final fun getHandlerCoroutineContext ()Lkotlin/coroutines/CoroutineContext;
|
||||
public final fun getTimeout-UwyO8pc ()J
|
||||
public final fun setEnforceStrictCapabilities (Z)V
|
||||
public final fun setHandlerCoroutineContext (Lkotlin/coroutines/CoroutineContext;)V
|
||||
public final fun setTimeout-LRDsOJo (J)V
|
||||
}
|
||||
|
||||
|
|
@ -107,8 +112,16 @@ public final class io/modelcontextprotocol/kotlin/sdk/shared/ReadBufferKt {
|
|||
public static final fun serializeMessage (Lio/modelcontextprotocol/kotlin/sdk/types/JSONRPCMessage;)Ljava/lang/String;
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra {
|
||||
public fun <init> ()V
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra : kotlin/coroutines/AbstractCoroutineContextElement {
|
||||
public static final field Key Lio/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra$Key;
|
||||
public final fun getMethod ()Lio/modelcontextprotocol/kotlin/sdk/types/Method;
|
||||
public final fun getRequestId ()Lio/modelcontextprotocol/kotlin/sdk/types/RequestId;
|
||||
public final fun sendNotification (Lio/modelcontextprotocol/kotlin/sdk/types/Notification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public final fun sendRequest (Lio/modelcontextprotocol/kotlin/sdk/types/Request;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public static synthetic fun sendRequest$default (Lio/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra;Lio/modelcontextprotocol/kotlin/sdk/types/Request;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestHandlerExtra$Key : kotlin/coroutines/CoroutineContext$Key {
|
||||
}
|
||||
|
||||
public final class io/modelcontextprotocol/kotlin/sdk/shared/RequestOptions : io/modelcontextprotocol/kotlin/sdk/shared/TransportSendOptions {
|
||||
|
|
|
|||
|
|
@ -27,15 +27,33 @@ import kotlinx.atomicfu.AtomicRef
|
|||
import kotlinx.atomicfu.atomic
|
||||
import kotlinx.atomicfu.getAndUpdate
|
||||
import kotlinx.atomicfu.update
|
||||
import kotlinx.collections.immutable.PersistentList
|
||||
import kotlinx.collections.immutable.PersistentMap
|
||||
import kotlinx.collections.immutable.persistentListOf
|
||||
import kotlinx.collections.immutable.persistentMapOf
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.CoroutineName
|
||||
import kotlinx.coroutines.CoroutineScope
|
||||
import kotlinx.coroutines.CoroutineStart
|
||||
import kotlinx.coroutines.Deferred
|
||||
import kotlinx.coroutines.TimeoutCancellationException
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.coroutines.Job
|
||||
import kotlinx.coroutines.NonCancellable
|
||||
import kotlinx.coroutines.SupervisorJob
|
||||
import kotlinx.coroutines.cancel
|
||||
import kotlinx.coroutines.currentCoroutineContext
|
||||
import kotlinx.coroutines.job
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.sync.Semaphore
|
||||
import kotlinx.coroutines.sync.withPermit
|
||||
import kotlinx.coroutines.withContext
|
||||
import kotlinx.coroutines.withTimeoutOrNull
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonPrimitive
|
||||
import kotlinx.serialization.json.encodeToJsonElement
|
||||
import kotlin.coroutines.AbstractCoroutineContextElement
|
||||
import kotlin.coroutines.ContinuationInterceptor
|
||||
import kotlin.coroutines.CoroutineContext
|
||||
import kotlin.coroutines.cancellation.CancellationException
|
||||
import kotlin.time.Duration
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
|
@ -53,17 +71,26 @@ public typealias ProgressCallback = (Progress) -> Unit
|
|||
/**
|
||||
* Additional initialization options.
|
||||
*
|
||||
* @property enforceStrictCapabilities whether to restrict emitted requests to only those that the
|
||||
* remote side has indicated it can handle through its advertised capabilities.
|
||||
* This does NOT affect checking of _local_ side capabilities, as it is considered a logic error
|
||||
* to mis-specify those.
|
||||
* Currently defaults to `false` for backwards compatibility with SDK versions that did not advertise
|
||||
* capabilities correctly; in the future, this will default to `true`.
|
||||
* [handlerCoroutineContext] is read once when [Protocol.connect] attaches a transport and stays
|
||||
* fixed until the next [Protocol.connect].
|
||||
*
|
||||
* @property enforceStrictCapabilities whether to restrict emitted requests to only those the remote
|
||||
* side advertised support for. Local-side capability checks are unaffected. Defaults to `false` for
|
||||
* backward compatibility with SDK versions that did not advertise capabilities correctly.
|
||||
* @property timeout default timeout for outgoing requests
|
||||
* @property handlerCoroutineContext coroutine context used to run inbound request and notification
|
||||
* handlers once initialization has completed. Must contain a real dispatching
|
||||
* [kotlin.coroutines.ContinuationInterceptor]. `Dispatchers.Unconfined` and unconfined test
|
||||
* dispatchers are unsupported, since handler resumptions would then run on the transport read loop
|
||||
* and reintroduce head-of-line blocking. Any [kotlinx.coroutines.Job] in this context is ignored.
|
||||
*/
|
||||
public open class ProtocolOptions(
|
||||
public var enforceStrictCapabilities: Boolean = false,
|
||||
public var timeout: Duration = DEFAULT_REQUEST_TIMEOUT,
|
||||
public var handlerCoroutineContext: CoroutineContext = Dispatchers.Default,
|
||||
// Internal handler-concurrency bounds, read once at connect and not part of the public surface.
|
||||
internal var maxConcurrentHandlers: Int = DEFAULT_MAX_CONCURRENT_HANDLERS,
|
||||
internal var maxInFlightHandlers: Int = DEFAULT_MAX_IN_FLIGHT_HANDLERS,
|
||||
)
|
||||
|
||||
/**
|
||||
|
|
@ -71,6 +98,12 @@ public open class ProtocolOptions(
|
|||
*/
|
||||
public val DEFAULT_REQUEST_TIMEOUT: Duration = 60.seconds
|
||||
|
||||
/** Default cap on concurrently executing inbound handlers per connection. */
|
||||
internal const val DEFAULT_MAX_CONCURRENT_HANDLERS: Int = 64
|
||||
|
||||
/** Default cap on launched-but-not-completed inbound handlers per connection. */
|
||||
internal const val DEFAULT_MAX_IN_FLIGHT_HANDLERS: Int = 256
|
||||
|
||||
/**
|
||||
* Options that can be given per request.
|
||||
*
|
||||
|
|
@ -121,7 +154,7 @@ public class RequestOptions(
|
|||
|
||||
override fun hashCode(): Int {
|
||||
var result = super.hashCode()
|
||||
result = 31 * result + (onProgress?.hashCode() ?: 0)
|
||||
result = 31 * result + onProgress.hashCode()
|
||||
result = 31 * result + timeout.hashCode()
|
||||
return result
|
||||
}
|
||||
|
|
@ -132,11 +165,70 @@ public class RequestOptions(
|
|||
|
||||
/**
|
||||
* Extra data given to request handlers.
|
||||
*
|
||||
* One instance exists per in-flight inbound request. It is passed as the second parameter to
|
||||
* low-level handlers registered via [Protocol.setRequestHandler] **and** installed into the
|
||||
* handler coroutine's [CoroutineContext], so code nested anywhere under the handler can reach it
|
||||
* via [currentRequestHandlerExtra] without parameter plumbing.
|
||||
*
|
||||
* @property requestId the JSON-RPC id of the request being handled
|
||||
* @property method the request's method
|
||||
*/
|
||||
public class RequestHandlerExtra
|
||||
public class RequestHandlerExtra internal constructor(
|
||||
public val requestId: RequestId,
|
||||
public val method: Method,
|
||||
private val protocol: Protocol,
|
||||
internal val capturedTransport: Transport,
|
||||
) : AbstractCoroutineContextElement(Key) {
|
||||
|
||||
/** Context key for retrieving the extra from a handler coroutine's context. */
|
||||
public companion object Key : CoroutineContext.Key<RequestHandlerExtra>
|
||||
|
||||
/**
|
||||
* Sends a notification tagged with `relatedRequestId = `[requestId], letting transports
|
||||
* associate it with the request being handled (e.g. route it onto the request's SSE stream).
|
||||
*/
|
||||
public suspend fun sendNotification(notification: Notification) {
|
||||
protocol.notification(notification, relatedRequestId = requestId)
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a request to the peer tagged with `relatedRequestId = `[requestId].
|
||||
*
|
||||
* Mirrors [Protocol.request]; cancellation of the handler coroutine propagates into this call.
|
||||
*/
|
||||
public suspend fun <T : RequestResult> sendRequest(request: Request, options: RequestOptions? = null): T {
|
||||
val base = options ?: RequestOptions()
|
||||
return protocol.request(
|
||||
request,
|
||||
RequestOptions(
|
||||
relatedRequestId = requestId,
|
||||
resumptionToken = base.resumptionToken,
|
||||
onResumptionToken = base.onResumptionToken,
|
||||
onProgress = base.onProgress,
|
||||
timeout = base.timeout,
|
||||
),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/** The [RequestHandlerExtra] of the MCP request being handled by the current coroutine, or `null` outside a handler. */
|
||||
public suspend fun currentRequestHandlerExtra(): RequestHandlerExtra? = currentCoroutineContext()[RequestHandlerExtra]
|
||||
|
||||
internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) }
|
||||
|
||||
/** Cap on how many recently-cancelled outbound request ids are remembered to quiet late responses/progress. */
|
||||
private const val CANCELLED_REQUEST_IDS_REMEMBERED = 256
|
||||
|
||||
// Bypass set: exempt from both bound tiers so the connection cannot self-block.
|
||||
// Responses are not listed — they are handled inline before dispatch.
|
||||
private val CONTROL_METHODS: Set<String> = setOf(
|
||||
Method.Defined.Ping.value,
|
||||
Method.Defined.NotificationsCancelled.value,
|
||||
Method.Defined.NotificationsProgress.value,
|
||||
Method.Defined.NotificationsInitialized.value,
|
||||
)
|
||||
|
||||
/**
|
||||
* Implements MCP protocol framing on top of a pluggable transport, including
|
||||
* features like request/response linking, notifications, and progress.
|
||||
|
|
@ -144,9 +236,11 @@ internal val COMPLETED = CompletableDeferred(Unit).also { it.complete(Unit) }
|
|||
* @property options protocol-level configuration; `null` falls back to defaults
|
||||
*/
|
||||
public abstract class Protocol(@PublishedApi internal val options: ProtocolOptions?) {
|
||||
private val connectionRef: AtomicRef<Connection?> = atomic(null)
|
||||
|
||||
/** The active transport, or `null` if not connected. */
|
||||
public var transport: Transport? = null
|
||||
private set
|
||||
public val transport: Transport?
|
||||
get() = connectionRef.value?.transport
|
||||
|
||||
private val _requestHandlers:
|
||||
AtomicRef<PersistentMap<String, suspend (JSONRPCRequest, RequestHandlerExtra) -> RequestResult?>> =
|
||||
|
|
@ -184,6 +278,21 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
public val progressHandlers: Map<ProgressToken, ProgressCallback>
|
||||
get() = _progressHandlers.value
|
||||
|
||||
private val recentlyCancelledRequestIds: AtomicRef<PersistentList<RequestId>> =
|
||||
atomic(persistentListOf())
|
||||
|
||||
/** Remembers a locally-cancelled outbound request id so its late response/progress is quieted. */
|
||||
private fun rememberCancelledRequestId(id: RequestId) {
|
||||
recentlyCancelledRequestIds.update { current ->
|
||||
val appended = current.adding(id)
|
||||
if (appended.size > CANCELLED_REQUEST_IDS_REMEMBERED) appended.removingAt(0) else appended
|
||||
}
|
||||
}
|
||||
|
||||
/** Whether [id] was recently cancelled locally (see [rememberCancelledRequestId]). */
|
||||
private fun isRecentlyCancelled(id: RequestId?): Boolean =
|
||||
id != null && recentlyCancelledRequestIds.value.contains(id)
|
||||
|
||||
/**
|
||||
* Callback for when the connection is closed for any reason.
|
||||
*
|
||||
|
|
@ -213,6 +322,11 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
public var fallbackNotificationHandler: (suspend (notification: JSONRPCNotification) -> Unit)? = null
|
||||
|
||||
init {
|
||||
setNotificationHandler<CancelledNotification>(Method.Defined.NotificationsCancelled) { notification ->
|
||||
handleCancelledNotification(notification)
|
||||
COMPLETED
|
||||
}
|
||||
|
||||
setNotificationHandler<ProgressNotification>(Method.Defined.NotificationsProgress) { notification ->
|
||||
onProgress(notification)
|
||||
COMPLETED
|
||||
|
|
@ -223,15 +337,64 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Enables concurrent dispatch of inbound messages on the current connection.
|
||||
*
|
||||
* Called by subclasses when the MCP initialization phase completes (the client after the
|
||||
* initialize exchange; the server session on receiving `notifications/initialized`). Until
|
||||
* then, inbound messages are handled inline in arrival order. Reset by [connect].
|
||||
*/
|
||||
protected fun enableConcurrentDispatch() {
|
||||
val connection = connectionRef.value ?: return
|
||||
connection.concurrentDispatchEnabled.value = true
|
||||
}
|
||||
|
||||
/**
|
||||
* Router hook invoked whenever `notifications/initialized` is received, before handler lookup.
|
||||
*
|
||||
* Runs regardless of which notification handler is registered for the method, so replacing the
|
||||
* handler via [setNotificationHandler] cannot disable subclass initialization logic.
|
||||
*/
|
||||
protected open fun onInitializedNotification() {}
|
||||
|
||||
/**
|
||||
* Attaches to the given transport, starts it, and starts listening for messages.
|
||||
*
|
||||
* The Protocol object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward.
|
||||
* 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.
|
||||
*
|
||||
* @throws IllegalStateException if this Protocol is already connected; call [close] first.
|
||||
*/
|
||||
public open suspend fun connect(transport: Transport) {
|
||||
this.transport = transport
|
||||
val maxConcurrent = options?.maxConcurrentHandlers ?: DEFAULT_MAX_CONCURRENT_HANDLERS
|
||||
val maxInFlight = options?.maxInFlightHandlers ?: DEFAULT_MAX_IN_FLIGHT_HANDLERS
|
||||
require(maxConcurrent > 0) { "maxConcurrentHandlers must be positive, but was $maxConcurrent" }
|
||||
require(maxInFlight >= maxConcurrent) {
|
||||
"maxInFlightHandlers ($maxInFlight) must be >= maxConcurrentHandlers ($maxConcurrent)"
|
||||
}
|
||||
val handlerContext = (options?.handlerCoroutineContext ?: Dispatchers.Default).minusKey(Job)
|
||||
val interceptor = handlerContext[ContinuationInterceptor]
|
||||
if (interceptor === Dispatchers.Unconfined ||
|
||||
interceptor?.let { it::class.simpleName?.contains("Unconfined") } == true
|
||||
) {
|
||||
logger.warn {
|
||||
"handlerCoroutineContext uses an unconfined dispatcher ($interceptor); " +
|
||||
"handler resumptions will run on the transport read loop and may block unrelated messages"
|
||||
}
|
||||
}
|
||||
val connection = Connection(
|
||||
transport = transport,
|
||||
handlerScope = CoroutineScope(SupervisorJob() + handlerContext + CoroutineName("McpProtocol")),
|
||||
executionSemaphore = Semaphore(maxConcurrent),
|
||||
maxInFlightHandlers = maxInFlight,
|
||||
)
|
||||
if (!connectionRef.compareAndSet(null, connection)) {
|
||||
connection.handlerScope.cancel()
|
||||
error("Protocol is already connected; close() the current connection before connecting a new transport")
|
||||
}
|
||||
|
||||
transport.onClose {
|
||||
doClose()
|
||||
doClose(connection)
|
||||
}
|
||||
|
||||
transport.onError {
|
||||
|
|
@ -239,24 +402,42 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
}
|
||||
|
||||
transport.onMessage { message ->
|
||||
if (connectionRef.value !== connection) {
|
||||
logger.trace { "Dropping message received after close: ${message::class.simpleName}" }
|
||||
return@onMessage
|
||||
}
|
||||
when (message) {
|
||||
is JSONRPCResponse -> onResponse(message, null)
|
||||
is JSONRPCRequest -> onRequest(message)
|
||||
is JSONRPCNotification -> onNotification(message)
|
||||
is JSONRPCRequest -> dispatchRequest(connection, message)
|
||||
is JSONRPCNotification -> dispatchNotification(connection, message)
|
||||
is JSONRPCError -> onResponse(null, message)
|
||||
is JSONRPCEmptyMessage -> Unit
|
||||
}
|
||||
}
|
||||
|
||||
logger.info { "Starting transport" }
|
||||
transport.start()
|
||||
try {
|
||||
transport.start()
|
||||
} catch (cause: Throwable) {
|
||||
// Roll back so the Protocol can be reconnected after a failed start; idempotent with
|
||||
// a transport-initiated doClose (both CAS on the same connection).
|
||||
if (connectionRef.compareAndSet(connection, null)) {
|
||||
connection.handlerScope.cancel()
|
||||
}
|
||||
throw cause
|
||||
}
|
||||
}
|
||||
|
||||
private fun doClose() {
|
||||
val handlersToNotify = _responseHandlers.value.values.toList()
|
||||
_responseHandlers.getAndSet(persistentMapOf())
|
||||
private fun doClose(connection: Connection) {
|
||||
// A stale onClose from a previous transport must not tear down the successor connection.
|
||||
if (!connectionRef.compareAndSet(connection, null)) {
|
||||
logger.trace { "Ignoring close signal from a stale transport" }
|
||||
return
|
||||
}
|
||||
val handlersToNotify = _responseHandlers.getAndSet(persistentMapOf()).values.toList()
|
||||
_progressHandlers.getAndSet(persistentMapOf())
|
||||
transport = null
|
||||
connection.inFlightRequestJobs.getAndSet(persistentMapOf())
|
||||
connection.handlerScope.cancel()
|
||||
onClose()
|
||||
|
||||
val error = McpException(RPCError.ErrorCode.CONNECTION_CLOSED, "Connection closed")
|
||||
|
|
@ -265,6 +446,110 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
}
|
||||
}
|
||||
|
||||
private suspend fun dispatchRequest(connection: Connection, request: JSONRPCRequest) {
|
||||
if (!connection.concurrentDispatchEnabled.value) {
|
||||
// Serial phase (before the peer completes MCP initialization): inline, in arrival
|
||||
// order — the delivering coroutine awaits the handler.
|
||||
onRequest(request, connection, trackCancellation = false)
|
||||
return
|
||||
}
|
||||
val bypass = request.method in CONTROL_METHODS
|
||||
if (!bypass && !tryAdmit(connection)) {
|
||||
rejectOverloadedRequest(connection, request)
|
||||
return
|
||||
}
|
||||
// UNDISPATCHED starts the handler synchronously in the delivering coroutine, preserving
|
||||
// arrival order, then moves to handlerCoroutineContext at its first suspension.
|
||||
connection.handlerScope.launch(start = CoroutineStart.UNDISPATCHED) {
|
||||
// Register before any suspension point so a following notifications/cancelled can find
|
||||
// and cancel even a handler still parked on admission.
|
||||
registerInFlight(connection, request.id, coroutineContext.job)
|
||||
try {
|
||||
if (bypass) {
|
||||
onRequest(request, connection, trackCancellation = true)
|
||||
} else {
|
||||
connection.executionSemaphore.withPermit {
|
||||
onRequest(request, connection, trackCancellation = true)
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
if (!bypass) connection.inFlightCount.decrementAndGet()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun dispatchNotification(connection: Connection, notification: JSONRPCNotification) {
|
||||
if (notification.method == Method.Defined.NotificationsInitialized.value) {
|
||||
onInitializedNotification()
|
||||
}
|
||||
if (!connection.concurrentDispatchEnabled.value) {
|
||||
onNotification(notification)
|
||||
return
|
||||
}
|
||||
val bypass = notification.method in CONTROL_METHODS
|
||||
if (!bypass && !tryAdmit(connection)) {
|
||||
val dropped = IllegalStateException(
|
||||
"Dropped notification ${notification.method}: too many in-flight messages",
|
||||
)
|
||||
logger.warn { dropped.message }
|
||||
onError(dropped)
|
||||
return
|
||||
}
|
||||
connection.handlerScope.launch(start = CoroutineStart.UNDISPATCHED) {
|
||||
try {
|
||||
if (bypass) {
|
||||
onNotification(notification)
|
||||
} else {
|
||||
connection.executionSemaphore.withPermit { onNotification(notification) }
|
||||
}
|
||||
} finally {
|
||||
if (!bypass) connection.inFlightCount.decrementAndGet()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private fun tryAdmit(connection: Connection): Boolean {
|
||||
while (true) {
|
||||
val current = connection.inFlightCount.value
|
||||
if (current >= connection.maxInFlightHandlers) return false
|
||||
if (connection.inFlightCount.compareAndSet(current, current + 1)) return true
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun rejectOverloadedRequest(connection: Connection, request: JSONRPCRequest) {
|
||||
logger.warn {
|
||||
"Rejecting request ${request.method} (id: ${request.id}): " +
|
||||
"in-flight handler limit (${connection.maxInFlightHandlers}) reached"
|
||||
}
|
||||
try {
|
||||
connection.transport.send(
|
||||
JSONRPCError(
|
||||
id = request.id,
|
||||
error = RPCError(
|
||||
code = RPCError.ErrorCode.INTERNAL_ERROR,
|
||||
message = "Server is busy: too many in-flight messages",
|
||||
),
|
||||
),
|
||||
)
|
||||
} catch (e: CancellationException) {
|
||||
throw e
|
||||
} catch (cause: Throwable) {
|
||||
logger.error(cause) { "Failed to send busy rejection for request id: ${request.id}" }
|
||||
onError(cause)
|
||||
}
|
||||
}
|
||||
|
||||
private fun handleCancelledNotification(notification: CancelledNotification) {
|
||||
val connection = connectionRef.value ?: return
|
||||
val requestId = notification.params.requestId
|
||||
val job = connection.inFlightRequestJobs.value[requestId]
|
||||
if (job == null) {
|
||||
logger.trace { "Ignoring cancellation for unknown or already-completed request: $requestId" }
|
||||
return
|
||||
}
|
||||
job.cancel("Cancelled by peer: ${notification.params.reason ?: "unknown"}")
|
||||
}
|
||||
|
||||
private suspend fun onNotification(notification: JSONRPCNotification) {
|
||||
logger.trace { "Received notification: ${notification.method}" }
|
||||
|
||||
|
|
@ -284,15 +569,16 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
}
|
||||
}
|
||||
|
||||
private suspend fun onRequest(request: JSONRPCRequest) {
|
||||
private suspend fun onRequest(request: JSONRPCRequest, connection: Connection, trackCancellation: Boolean) {
|
||||
logger.trace { "Received request: ${request.method} (id: ${request.id})" }
|
||||
|
||||
val handler = requestHandlers[request.method] ?: fallbackRequestHandler
|
||||
val capturedTransport = connection.transport
|
||||
|
||||
if (handler === null) {
|
||||
logger.trace { "No handler found for request: ${request.method}" }
|
||||
try {
|
||||
transport?.send(
|
||||
capturedTransport.send(
|
||||
JSONRPCError(
|
||||
id = request.id,
|
||||
error = RPCError(
|
||||
|
|
@ -310,36 +596,77 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
return
|
||||
}
|
||||
|
||||
try {
|
||||
val result = handler(request, RequestHandlerExtra())
|
||||
logger.trace { "Request handled successfully: ${request.method} (id: ${request.id})" }
|
||||
val extra = RequestHandlerExtra(
|
||||
requestId = request.id,
|
||||
method = Method.from(request.method),
|
||||
protocol = this,
|
||||
capturedTransport = capturedTransport,
|
||||
)
|
||||
|
||||
transport?.send(
|
||||
JSONRPCResponse(
|
||||
id = request.id,
|
||||
result = result ?: EmptyResult(),
|
||||
),
|
||||
)
|
||||
// Registration in the in-flight registry already happened in dispatchRequest's launch
|
||||
// block (synchronously, pre-suspension). Here we only need the job for suppression checks.
|
||||
val handlerJob = if (trackCancellation) currentCoroutineContext().job else null
|
||||
|
||||
withContext(extra) {
|
||||
try {
|
||||
val result = handler(request, extra)
|
||||
logger.trace { "Request handled successfully: ${request.method} (id: ${request.id})" }
|
||||
if (handlerJob?.isCancelled == true) {
|
||||
logger.trace { "Suppressing response for request cancelled by peer (id: ${request.id})" }
|
||||
return@withContext
|
||||
}
|
||||
capturedTransport.send(
|
||||
JSONRPCResponse(
|
||||
id = request.id,
|
||||
result = result ?: EmptyResult(),
|
||||
),
|
||||
)
|
||||
} catch (e: CancellationException) {
|
||||
// Intentional exception to the always-rethrow-CE convention, only at this boundary.
|
||||
if (handlerJob == null) throw e // serial phase: inline dispatch is not response-suppressed
|
||||
if (handlerJob.isCancelled) throw e // genuine peer or close cancel: suppress the response
|
||||
// CE escaped a live handler job (e.g. a leaked inner withTimeout). Answer
|
||||
// INTERNAL_ERROR so the peer does not hang until its own timeout.
|
||||
respondWithError(capturedTransport, request, e)
|
||||
} catch (cause: Throwable) {
|
||||
if (handlerJob?.isCancelled == true) {
|
||||
logger.trace { "Suppressing error response for request cancelled by peer (id: ${request.id})" }
|
||||
return@withContext
|
||||
}
|
||||
respondWithError(capturedTransport, request, cause)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private fun registerInFlight(connection: Connection, id: RequestId, job: Job) {
|
||||
val previous = connection.inFlightRequestJobs.getAndUpdate { current -> current.putting(id, job) }[id]
|
||||
if (previous != null) {
|
||||
logger.warn { "Duplicate in-flight request id $id; replacing the previous handler job" }
|
||||
}
|
||||
job.invokeOnCompletion {
|
||||
// Identity guard: never evict a successor job registered under a reused id.
|
||||
connection.inFlightRequestJobs.update { current ->
|
||||
if (current[id] === job) current.removing(id) else current
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun respondWithError(transport: Transport, request: JSONRPCRequest, cause: Throwable) {
|
||||
logger.error(cause) { "Error handling request: ${request.method} (id: ${request.id})" }
|
||||
try {
|
||||
val rpcError = if (cause is McpException) {
|
||||
RPCError(code = cause.code, message = cause.message.orEmpty(), data = cause.data)
|
||||
} else {
|
||||
RPCError(code = RPCError.ErrorCode.INTERNAL_ERROR, message = cause.message ?: "Internal error")
|
||||
}
|
||||
transport.send(JSONRPCError(id = request.id, error = rpcError))
|
||||
} catch (e: CancellationException) {
|
||||
throw e
|
||||
} catch (cause: Throwable) {
|
||||
logger.error(cause) { "Error handling request: ${request.method} (id: ${request.id})" }
|
||||
|
||||
try {
|
||||
val rpcError = if (cause is McpException) {
|
||||
RPCError(code = cause.code, message = cause.message.orEmpty(), data = cause.data)
|
||||
} else {
|
||||
RPCError(code = RPCError.ErrorCode.INTERNAL_ERROR, message = cause.message ?: "Internal error")
|
||||
}
|
||||
transport?.send(JSONRPCError(id = request.id, error = rpcError))
|
||||
} catch (e: CancellationException) {
|
||||
throw e
|
||||
} catch (sendError: Throwable) {
|
||||
logger.error(sendError) {
|
||||
"Failed to send error response for request: ${request.method} (id: ${request.id})"
|
||||
}
|
||||
// Optionally implement fallback behavior here
|
||||
} catch (sendError: Throwable) {
|
||||
logger.error(sendError) {
|
||||
"Failed to send error response for request: ${request.method} (id: ${request.id})"
|
||||
}
|
||||
onError(sendError)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -354,6 +681,10 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
|
||||
val handler = _progressHandlers.value[progressToken]
|
||||
if (handler == null) {
|
||||
if (isRecentlyCancelled(progressToken)) {
|
||||
logger.trace { "Ignoring progress for a locally cancelled request: $progressToken" }
|
||||
return
|
||||
}
|
||||
val error = Error(
|
||||
"Received a progress notification for an unknown token: ${McpJson.encodeToString(notification)}",
|
||||
)
|
||||
|
|
@ -370,7 +701,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
|
||||
val oldResponseHandlers = _responseHandlers.getAndUpdate { current ->
|
||||
if (messageId != null && messageId in current) {
|
||||
current.remove(messageId)
|
||||
current.removing(messageId)
|
||||
} else {
|
||||
current
|
||||
}
|
||||
|
|
@ -379,8 +710,12 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
val handler = oldResponseHandlers[messageId]
|
||||
|
||||
if (handler != null) {
|
||||
messageId?.let { msg -> _progressHandlers.update { it.remove(msg) } }
|
||||
messageId?.let { msg -> _progressHandlers.update { it.removing(msg) } }
|
||||
} else {
|
||||
if (isRecentlyCancelled(messageId)) {
|
||||
logger.trace { "Ignoring response for a locally cancelled request: $messageId" }
|
||||
return
|
||||
}
|
||||
onError(
|
||||
IllegalStateException(
|
||||
"Received a response for an unknown message ID: ${McpJson.encodeToString(error ?: response)}",
|
||||
|
|
@ -438,7 +773,8 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
public suspend fun <T : RequestResult> request(request: Request, options: RequestOptions? = null): T {
|
||||
logger.trace { "Sending request: ${request.method}" }
|
||||
val result = CompletableDeferred<T>()
|
||||
val transport = transport ?: error("Not connected")
|
||||
val connection = connectionRef.value ?: error("Not connected")
|
||||
val transport = connection.transport
|
||||
|
||||
if (this@Protocol.options?.enforceStrictCapabilities == true) {
|
||||
assertCapabilityForMethod(request.method)
|
||||
|
|
@ -448,7 +784,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
options?.onProgress?.let { progressHandler ->
|
||||
logger.trace { "Registering progress handler for request id: $id" }
|
||||
_progressHandlers.update { current ->
|
||||
current.put(id, progressHandler)
|
||||
current.putting(id, progressHandler)
|
||||
}
|
||||
|
||||
val paramsObject = (this.params as? JsonObject) ?: JsonObject(emptyMap())
|
||||
|
|
@ -467,10 +803,10 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
val jsonRpcRequestId = jsonRpcRequest.id
|
||||
|
||||
_responseHandlers.update { current ->
|
||||
current.put(jsonRpcRequestId) { response, error ->
|
||||
current.putting(jsonRpcRequestId) { response, error ->
|
||||
if (error != null) {
|
||||
result.completeExceptionally(error)
|
||||
return@put
|
||||
return@putting
|
||||
}
|
||||
|
||||
try {
|
||||
|
|
@ -482,42 +818,84 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
}
|
||||
}
|
||||
|
||||
val cancel: suspend (Throwable) -> Unit = { reason: Throwable ->
|
||||
_responseHandlers.update { current -> current.remove(jsonRpcRequestId) }
|
||||
_progressHandlers.update { current -> current.remove(jsonRpcRequestId) }
|
||||
if (connectionRef.value !== connection) {
|
||||
// Connection closed or replaced during registration; nothing would complete this handler.
|
||||
_responseHandlers.update { it.removing(jsonRpcRequestId) }
|
||||
_progressHandlers.update { it.removing(jsonRpcRequestId) }
|
||||
throw McpException(RPCError.ErrorCode.CONNECTION_CLOSED, "Connection closed")
|
||||
}
|
||||
|
||||
val notification = CancelledNotification(
|
||||
params = CancelledNotificationParams(
|
||||
requestId = jsonRpcRequestId,
|
||||
reason = reason.message ?: "Unknown",
|
||||
),
|
||||
)
|
||||
val cancelPending: suspend (reason: Throwable, notifyPeer: Boolean) -> Unit = { reason, notifyPeer ->
|
||||
// Remember before removing: a racing onResponse/onProgress that sees the handler gone
|
||||
// must also see the id, or it reports a late message as unknown. Order is load-bearing.
|
||||
rememberCancelledRequestId(jsonRpcRequestId)
|
||||
_responseHandlers.update { current -> current.removing(jsonRpcRequestId) }
|
||||
_progressHandlers.update { current -> current.removing(jsonRpcRequestId) }
|
||||
|
||||
val jsonRpcNotification = notification.toJSON()
|
||||
|
||||
transport.send(jsonRpcNotification, options)
|
||||
if (notifyPeer) {
|
||||
val notification = CancelledNotification(
|
||||
params = CancelledNotificationParams(
|
||||
requestId = jsonRpcRequestId,
|
||||
reason = reason.message ?: "Unknown",
|
||||
),
|
||||
)
|
||||
transport.send(notification.toJSON(), options)
|
||||
}
|
||||
|
||||
result.completeExceptionally(reason)
|
||||
}
|
||||
|
||||
// The MCP spec forbids cancelling `initialize`; local cleanup still runs,
|
||||
// but no notifications/cancelled goes on the wire (timeout and cancel paths alike).
|
||||
val notifyPeerOnCancel = request.method != Method.Defined.Initialize
|
||||
|
||||
val timeout = options?.timeout ?: DEFAULT_REQUEST_TIMEOUT
|
||||
try {
|
||||
withTimeout(timeout) {
|
||||
val response = withTimeoutOrNull(timeout) {
|
||||
logger.trace { "Sending request message with id: $jsonRpcRequestId" }
|
||||
this@Protocol.transport?.send(jsonRpcRequest, options)
|
||||
transport.send(jsonRpcRequest, options)
|
||||
result.await()
|
||||
}
|
||||
return result.await()
|
||||
} catch (cause: TimeoutCancellationException) {
|
||||
logger.error { "Request timed out after ${timeout.inWholeMilliseconds}ms: ${request.method}" }
|
||||
cancel(
|
||||
McpException(
|
||||
if (response == null) {
|
||||
// Our own request timeout expired. An outer withTimeout's cancellation propagates
|
||||
// through withTimeoutOrNull instead of returning null and is handled by the catch below.
|
||||
logger.error { "Request timed out after ${timeout.inWholeMilliseconds}ms: ${request.method}" }
|
||||
val timeoutError = McpException(
|
||||
code = RPCError.ErrorCode.REQUEST_TIMEOUT,
|
||||
message = "Request timed out",
|
||||
data = JsonObject(mutableMapOf("timeout" to JsonPrimitive(timeout.inWholeMilliseconds))),
|
||||
),
|
||||
)
|
||||
result.cancel(cause)
|
||||
)
|
||||
withContext(NonCancellable) {
|
||||
try {
|
||||
cancelPending(timeoutError, notifyPeerOnCancel)
|
||||
} catch (e: Throwable) {
|
||||
logger.warn(e) { "Failed to notify peer about timed-out request" }
|
||||
onError(e)
|
||||
}
|
||||
}
|
||||
result.cancel()
|
||||
throw timeoutError
|
||||
}
|
||||
return response
|
||||
} catch (cause: CancellationException) {
|
||||
// Plain caller cancellation, or an outer withTimeout that propagated through
|
||||
// withTimeoutOrNull (only our own deadline returns null there). NonCancellable so the
|
||||
// cancellation notification still goes out even though the caller's job is cancelled.
|
||||
// A secondary send failure must never replace the original CancellationException.
|
||||
withContext(NonCancellable) {
|
||||
try {
|
||||
cancelPending(cause, notifyPeerOnCancel)
|
||||
} catch (e: Throwable) {
|
||||
logger.warn(e) { "Failed to notify peer about cancelled request" }
|
||||
onError(e)
|
||||
}
|
||||
}
|
||||
throw cause
|
||||
} finally {
|
||||
// Backstop for the failed-send path. The success, timeout and cancellation paths already
|
||||
// removed these handlers, so this only bites when the outbound send failed.
|
||||
_responseHandlers.update { it.removing(jsonRpcRequestId) }
|
||||
_progressHandlers.update { it.removing(jsonRpcRequestId) }
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -538,6 +916,11 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
* Registers a handler to invoke when this protocol object receives a request with the given method.
|
||||
*
|
||||
* Note that this will replace any previous request handler for the same method.
|
||||
*
|
||||
* Handlers may run concurrently after initialization and are cancelled cooperatively, so always
|
||||
* re-throw [CancellationException]. Replacing a built-in control-method handler (`ping`,
|
||||
* `notifications/cancelled`, `notifications/progress`, `notifications/initialized`) with slow
|
||||
* code degrades connection liveness.
|
||||
*/
|
||||
public inline fun <reified T : Request> setRequestHandler(
|
||||
method: Method,
|
||||
|
|
@ -556,7 +939,7 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
val wrapped = wrapRequestHandler(method, block)
|
||||
|
||||
_requestHandlers.update { current ->
|
||||
current.put(method.value) { jSONRPCRequest, extraHandler ->
|
||||
current.putting(method.value) { jSONRPCRequest, extraHandler ->
|
||||
val request = jSONRPCRequest.fromJSON()
|
||||
val response = wrapped(request as T, extraHandler)
|
||||
response
|
||||
|
|
@ -582,17 +965,23 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
* Removes the request handler for the given method.
|
||||
*/
|
||||
public fun removeRequestHandler(method: Method) {
|
||||
_requestHandlers.update { current -> current.remove(method.value) }
|
||||
_requestHandlers.update { current -> current.removing(method.value) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a handler to invoke when this protocol object receives a notification with the given method.
|
||||
*
|
||||
* Note that this will replace any previous notification handler for the same method.
|
||||
*
|
||||
* Handlers may run concurrently after initialization and are cancelled cooperatively, so always
|
||||
* re-throw [CancellationException]. Replacing a built-in control-method handler (`ping`,
|
||||
* `notifications/cancelled`, `notifications/progress`, `notifications/initialized`) with slow
|
||||
* code degrades connection liveness. Replacing `notifications/cancelled` in particular disables
|
||||
* inbound request cancellation entirely, since that handler is what cancels the running job.
|
||||
*/
|
||||
public fun <T : Notification> setNotificationHandler(method: Method, handler: (notification: T) -> Deferred<Unit>) {
|
||||
_notificationHandlers.update { current ->
|
||||
current.put(method.value) {
|
||||
current.putting(method.value) {
|
||||
@Suppress("UNCHECKED_CAST")
|
||||
handler(it.fromJSON() as T)
|
||||
}
|
||||
|
|
@ -603,6 +992,22 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
|
|||
* Removes the notification handler for the given method.
|
||||
*/
|
||||
public fun removeNotificationHandler(method: Method) {
|
||||
_notificationHandlers.update { current -> current.remove(method.value) }
|
||||
_notificationHandlers.update { current -> current.removing(method.value) }
|
||||
}
|
||||
|
||||
private class Connection(
|
||||
val transport: Transport,
|
||||
val handlerScope: CoroutineScope,
|
||||
val executionSemaphore: Semaphore,
|
||||
val maxInFlightHandlers: Int,
|
||||
) {
|
||||
/** Init gate: inbound dispatch is inline/serial until this flips. */
|
||||
val concurrentDispatchEnabled = atomic(false)
|
||||
|
||||
/** Admission tier: launched-but-not-completed handler jobs (running + parked). */
|
||||
val inFlightCount = atomic(0)
|
||||
|
||||
/** In-flight request handler jobs, for `notifications/cancelled`. */
|
||||
val inFlightRequestJobs: AtomicRef<PersistentMap<RequestId, Job>> = atomic(persistentMapOf())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -53,4 +53,13 @@ public sealed interface Method {
|
|||
*/
|
||||
@Serializable
|
||||
public data class Custom(override val value: String) : Method
|
||||
|
||||
public companion object {
|
||||
private val definedByValue: Map<String, Defined> by lazy {
|
||||
Defined.entries.associateBy { it.value }
|
||||
}
|
||||
|
||||
/** Resolves a wire method string to a [Defined] entry when known, [Custom] otherwise. */
|
||||
internal fun from(value: String): Method = definedByValue[value] ?: Custom(value)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -59,18 +59,11 @@ internal object MethodSerializer : KSerializer<Method> {
|
|||
override val descriptor: SerialDescriptor =
|
||||
PrimitiveSerialDescriptor("io.modelcontextprotocol.kotlin.sdk.types.Method", PrimitiveKind.STRING)
|
||||
|
||||
private val definedMethods: Map<String, Method.Defined> by lazy {
|
||||
Method.Defined.entries.associateBy { it.value }
|
||||
}
|
||||
|
||||
override fun serialize(encoder: Encoder, value: Method) {
|
||||
encoder.encodeString(value.value)
|
||||
}
|
||||
|
||||
override fun deserialize(decoder: Decoder): Method {
|
||||
val decodedString = decoder.decodeString()
|
||||
return definedMethods[decodedString] ?: Method.Custom(decodedString)
|
||||
}
|
||||
override fun deserialize(decoder: Decoder): Method = Method.from(decoder.decodeString())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
|
|
|
|||
|
|
@ -0,0 +1,524 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
import io.kotest.assertions.throwables.shouldThrow
|
||||
import io.kotest.matchers.collections.shouldHaveSize
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.kotest.matchers.string.shouldContain
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification
|
||||
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.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.PingRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.async
|
||||
import kotlinx.coroutines.delay
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.test.StandardTestDispatcher
|
||||
import kotlinx.coroutines.test.TestScope
|
||||
import kotlinx.coroutines.test.advanceTimeBy
|
||||
import kotlinx.coroutines.test.advanceUntilIdle
|
||||
import kotlinx.coroutines.test.runCurrent
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.buildJsonObject
|
||||
import kotlinx.serialization.json.int
|
||||
import kotlinx.serialization.json.jsonPrimitive
|
||||
import kotlinx.serialization.json.put
|
||||
import kotlin.coroutines.cancellation.CancellationException
|
||||
import kotlin.test.Test
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
class ProtocolConcurrencyTest {
|
||||
|
||||
private fun TestScope.newProtocol(): TestProtocol =
|
||||
TestProtocol(ProtocolOptions(handlerCoroutineContext = StandardTestDispatcher(testScheduler)))
|
||||
|
||||
@Test
|
||||
fun `a suspended slow handler does not block a later fast request`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/slow") delay(10.seconds)
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow"))
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast"))
|
||||
runCurrent()
|
||||
|
||||
responsesOn(transport).map { it.id } shouldBe listOf(RequestId(2L))
|
||||
|
||||
advanceTimeBy(11.seconds)
|
||||
runCurrent()
|
||||
responsesOn(transport).map { it.id } shouldBe listOf(RequestId(2L), RequestId(1L))
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `dispatch is serial until concurrent dispatch is enabled`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
// gate NOT flipped
|
||||
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/slow") delay(10.seconds)
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
val delivery = launch {
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow"))
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast"))
|
||||
}
|
||||
runCurrent()
|
||||
// inline handling: the delivering coroutine is parked inside the slow handler
|
||||
delivery.isCompleted shouldBe false
|
||||
responsesOn(transport) shouldBe emptyList()
|
||||
|
||||
advanceTimeBy(11.seconds)
|
||||
runCurrent()
|
||||
delivery.isCompleted shouldBe true
|
||||
responsesOn(transport).map { it.id } shouldBe listOf(RequestId(1L), RequestId(2L))
|
||||
}
|
||||
|
||||
// the hook fires before handler lookup and cannot be disabled by a user handler
|
||||
@Test
|
||||
fun `initialized notification invokes the router hook even when a user handler is registered`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
var userHandlerCalled = false
|
||||
protocol.setNotificationHandler<InitializedNotification>(Method.Defined.NotificationsInitialized) {
|
||||
userHandlerCalled = true
|
||||
COMPLETED
|
||||
}
|
||||
transport.deliver(JSONRPCNotification(method = Method.Defined.NotificationsInitialized.value))
|
||||
runCurrent()
|
||||
|
||||
protocol.initializedNotificationCount shouldBe 1
|
||||
userHandlerCalled shouldBe true
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `cancelled notification cancels the in-flight handler and suppresses the response`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val entered = CompletableDeferred<Unit>()
|
||||
val sawCancellation = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
entered.complete(Unit)
|
||||
try {
|
||||
delay(10.seconds)
|
||||
} catch (e: CancellationException) {
|
||||
sawCancellation.complete(Unit)
|
||||
throw e
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(5L), method = "test/slow"))
|
||||
runCurrent()
|
||||
entered.await()
|
||||
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(5L), reason = "user gave up"))
|
||||
runCurrent()
|
||||
sawCancellation.await()
|
||||
advanceUntilIdle()
|
||||
|
||||
responsesOn(transport) shouldBe emptyList()
|
||||
errorsOn(transport) shouldBe emptyList() // JSONRPCError messages on the wire
|
||||
protocol.errors shouldBe emptyList()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `cancelled notification for an unknown id is a silent no-op`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(404L), reason = "nope"))
|
||||
advanceUntilIdle()
|
||||
|
||||
protocol.errors shouldBe emptyList()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `handler leaking CancellationException without being cancelled produces INTERNAL_ERROR`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
protocol.fallbackRequestHandler = { _, _ -> throw CancellationException("leaked") }
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(6L), method = "test/leak"))
|
||||
advanceUntilIdle()
|
||||
|
||||
val error = errorsOn(transport).single()
|
||||
error.id shouldBe RequestId(6L)
|
||||
error.error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR
|
||||
}
|
||||
|
||||
// suppression race variants
|
||||
@Test
|
||||
fun `handler that swallows cancellation and returns normally still gets no response`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val entered = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
entered.complete(Unit)
|
||||
try {
|
||||
delay(10.seconds)
|
||||
} catch (_: CancellationException) {
|
||||
// swallow (anti-pattern) and return normally
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(7L), method = "test/swallow"))
|
||||
runCurrent()
|
||||
entered.await()
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(7L), reason = "stop"))
|
||||
advanceUntilIdle()
|
||||
|
||||
responsesOn(transport) shouldBe emptyList()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `handler throwing non-CE while cancelled gets no error response`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val entered = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
entered.complete(Unit)
|
||||
try {
|
||||
delay(10.seconds)
|
||||
} catch (_: CancellationException) {
|
||||
throw IllegalStateException("boom")
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(8L), method = "test/throw"))
|
||||
runCurrent()
|
||||
entered.await()
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(8L), reason = "stop"))
|
||||
advanceUntilIdle()
|
||||
|
||||
errorsOn(transport) shouldBe emptyList()
|
||||
protocol.errors shouldBe emptyList()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `suspension-free notification handlers observe strict arrival order`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val seen = mutableListOf<Int>()
|
||||
protocol.fallbackNotificationHandler = { notification ->
|
||||
seen.add((notification.params as JsonObject).getValue("i").jsonPrimitive.int)
|
||||
}
|
||||
|
||||
repeat(20) { i ->
|
||||
transport.deliver(
|
||||
JSONRPCNotification(method = "test/event", params = buildJsonObject { put("i", i) }),
|
||||
)
|
||||
}
|
||||
advanceUntilIdle()
|
||||
|
||||
seen shouldBe (0 until 20).toList()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `reconnect resets the dispatch gate to serial`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport1 = RecordingTransport()
|
||||
protocol.connect(transport1)
|
||||
protocol.enableConcurrency()
|
||||
protocol.close()
|
||||
|
||||
val transport2 = RecordingTransport()
|
||||
protocol.connect(transport2)
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/slow") delay(10.seconds)
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
val delivery = launch {
|
||||
transport2.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow"))
|
||||
transport2.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast"))
|
||||
}
|
||||
runCurrent()
|
||||
delivery.isCompleted shouldBe false // serial again: gate reset
|
||||
advanceTimeBy(11.seconds)
|
||||
runCurrent()
|
||||
responsesOn(transport2).map { it.id } shouldBe listOf(RequestId(1L), RequestId(2L))
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `close cancels in-flight handlers suppresses responses and fails pending outbound requests`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val entered = CompletableDeferred<Unit>()
|
||||
val sawCancellation = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
entered.complete(Unit)
|
||||
try {
|
||||
delay(10.seconds)
|
||||
} catch (e: CancellationException) {
|
||||
sawCancellation.complete(Unit)
|
||||
throw e
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(9L), method = "test/slow"))
|
||||
runCurrent()
|
||||
entered.await()
|
||||
|
||||
// shouldThrow inside the async so its failure does not cancel the (non-supervisor) test scope.
|
||||
val pending = async { shouldThrow<McpException> { protocol.request<EmptyResult>(PingRequest()) } }
|
||||
transport.awaitRequest() // outbound ping on the wire
|
||||
|
||||
protocol.close()
|
||||
runCurrent()
|
||||
|
||||
sawCancellation.await()
|
||||
responsesOn(transport) shouldBe emptyList()
|
||||
val thrown = pending.await()
|
||||
thrown.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `message delivered after close is dropped without crash or response`() = runTest {
|
||||
val protocol = newProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
protocol.close()
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/anything"))
|
||||
advanceUntilIdle()
|
||||
|
||||
responsesOn(transport) shouldBe emptyList()
|
||||
errorsOn(transport) shouldBe emptyList()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `at most maxConcurrentHandlers handlers run concurrently while control messages bypass`() = runTest {
|
||||
val protocol = TestProtocol(
|
||||
ProtocolOptions(
|
||||
handlerCoroutineContext = StandardTestDispatcher(testScheduler),
|
||||
maxConcurrentHandlers = 2,
|
||||
),
|
||||
)
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
var entered = 0
|
||||
val release = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
entered += 1
|
||||
release.await()
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
repeat(4) { i -> transport.deliver(JSONRPCRequest(id = RequestId(i.toLong()), method = "test/slow")) }
|
||||
runCurrent()
|
||||
entered shouldBe 2 // exactly two running, two parked on the semaphore
|
||||
|
||||
// ping (bypass) is answered while saturated
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(100L), method = Method.Defined.Ping.value))
|
||||
runCurrent()
|
||||
responsesOn(transport).map { it.id } shouldBe listOf(RequestId(100L))
|
||||
|
||||
// responses (bypass, inline) are processed while saturated
|
||||
val pending = async { protocol.request<EmptyResult>(PingRequest()) }
|
||||
runCurrent()
|
||||
val outbound = transport.awaitRequest()
|
||||
transport.deliver(JSONRPCResponse(id = outbound.id, result = EmptyResult()))
|
||||
runCurrent()
|
||||
pending.isCompleted shouldBe true
|
||||
|
||||
release.complete(Unit)
|
||||
advanceUntilIdle()
|
||||
entered shouldBe 4
|
||||
responsesOn(transport) shouldHaveSize 5
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `flood beyond maxInFlightHandlers is rejected fail-fast and the read loop never suspends`() = runTest {
|
||||
val protocol = TestProtocol(
|
||||
ProtocolOptions(
|
||||
handlerCoroutineContext = StandardTestDispatcher(testScheduler),
|
||||
maxConcurrentHandlers = 1,
|
||||
maxInFlightHandlers = 4,
|
||||
),
|
||||
)
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val release = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
release.await()
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
// 10 requests delivered back-to-back; delivery must complete without advancing
|
||||
// virtual time — i.e. the read loop was never suspended on admission.
|
||||
val delivery = launch {
|
||||
repeat(10) { i -> transport.deliver(JSONRPCRequest(id = RequestId(i.toLong()), method = "test/slow")) }
|
||||
}
|
||||
runCurrent()
|
||||
delivery.isCompleted shouldBe true
|
||||
|
||||
// 4 admitted (1 running + 3 parked), 6 rejected immediately
|
||||
val busyErrors = errorsOn(transport)
|
||||
busyErrors shouldHaveSize 6
|
||||
busyErrors.forEach {
|
||||
it.error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR
|
||||
it.error.message shouldBe "Server is busy: too many in-flight messages"
|
||||
}
|
||||
|
||||
// ping still answered at full saturation
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(100L), method = Method.Defined.Ping.value))
|
||||
runCurrent()
|
||||
responsesOn(transport).map { it.id } shouldBe listOf(RequestId(100L))
|
||||
|
||||
release.complete(Unit)
|
||||
advanceUntilIdle()
|
||||
responsesOn(transport) shouldHaveSize 5 // ping + the 4 admitted
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `overflowing notifications are dropped and reported via onError`() = runTest {
|
||||
val protocol = TestProtocol(
|
||||
ProtocolOptions(
|
||||
handlerCoroutineContext = StandardTestDispatcher(testScheduler),
|
||||
maxConcurrentHandlers = 1,
|
||||
maxInFlightHandlers = 2,
|
||||
),
|
||||
)
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
// parks one handler with a request: holds the sole execution permit and one in-flight slot
|
||||
val release = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { _, _ ->
|
||||
release.await()
|
||||
EmptyResult()
|
||||
}
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/hog"))
|
||||
runCurrent()
|
||||
|
||||
// notification 1 is admitted (parked on the semaphore, behind the hogged permit);
|
||||
// notification 2 overflows the admission cap and is dropped instead of parked.
|
||||
val notificationRelease = CompletableDeferred<Unit>()
|
||||
protocol.fallbackNotificationHandler = { notificationRelease.await() }
|
||||
transport.deliver(JSONRPCNotification(method = "test/event"))
|
||||
transport.deliver(JSONRPCNotification(method = "test/event"))
|
||||
runCurrent()
|
||||
|
||||
protocol.errors shouldHaveSize 1
|
||||
protocol.errors.single().message.orEmpty() shouldContain "too many in-flight messages"
|
||||
errorsOn(transport) shouldBe emptyList() // no wire error for dropped notifications
|
||||
|
||||
release.complete(Unit)
|
||||
notificationRelease.complete(Unit)
|
||||
advanceUntilIdle()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `cancelling saturating handlers releases permits and lets parked handlers run`() = runTest {
|
||||
val protocol = TestProtocol(
|
||||
ProtocolOptions(
|
||||
handlerCoroutineContext = StandardTestDispatcher(testScheduler),
|
||||
maxConcurrentHandlers = 1,
|
||||
),
|
||||
)
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val secondCompleted = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/hog") {
|
||||
delay(1000.seconds) // cancellable park
|
||||
error("unreachable")
|
||||
}
|
||||
secondCompleted.complete(Unit)
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/hog")) // holds the permit
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/next")) // parked on the semaphore
|
||||
runCurrent()
|
||||
secondCompleted.isCompleted shouldBe false
|
||||
|
||||
// cancel the hog THROUGH the bypass while saturated
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(1L), reason = "make room"))
|
||||
advanceUntilIdle()
|
||||
|
||||
secondCompleted.isCompleted shouldBe true
|
||||
responsesOn(transport).map { it.id } shouldBe listOf(RequestId(2L)) // hog suppressed, next answered
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `a parked handler can itself be cancelled before it ever runs`() = runTest {
|
||||
val protocol = TestProtocol(
|
||||
ProtocolOptions(
|
||||
handlerCoroutineContext = StandardTestDispatcher(testScheduler),
|
||||
maxConcurrentHandlers = 1,
|
||||
),
|
||||
)
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/hog") {
|
||||
delay(1000.seconds) // cancellable park
|
||||
error("unreachable")
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/hog")) // holds the permit
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/next")) // parked on the semaphore
|
||||
runCurrent()
|
||||
|
||||
// cancel the parked handler before it ever acquires the permit, then cancel the hog too
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(2L), reason = "give up while parked"))
|
||||
transport.deliver(cancelledNotification(requestId = RequestId(1L), reason = "make room"))
|
||||
advanceUntilIdle()
|
||||
|
||||
// no response for id=2 at all: its withPermit acquire was cancelled before onRequest ever ran
|
||||
responsesOn(transport) shouldBe emptyList()
|
||||
errorsOn(transport) shouldBe emptyList()
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
import io.kotest.matchers.shouldBe
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlin.test.Test
|
||||
|
||||
class ProtocolOptionsTest {
|
||||
|
||||
@Test
|
||||
fun `defaults expose concurrency knobs`() {
|
||||
val options = ProtocolOptions()
|
||||
options.handlerCoroutineContext shouldBe Dispatchers.Default
|
||||
options.maxConcurrentHandlers shouldBe 64
|
||||
options.maxInFlightHandlers shouldBe 256
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,167 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
import io.kotest.assertions.throwables.shouldThrow
|
||||
import io.kotest.matchers.collections.shouldHaveSize
|
||||
import io.kotest.matchers.nulls.shouldNotBeNull
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.kotest.matchers.string.shouldContain
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.InitializeRequestParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.LATEST_PROTOCOL_VERSION
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpException
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpJson
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.PingRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.TimeoutCancellationException
|
||||
import kotlinx.coroutines.async
|
||||
import kotlinx.coroutines.cancelAndJoin
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.test.advanceTimeBy
|
||||
import kotlinx.coroutines.test.runCurrent
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import kotlinx.serialization.json.decodeFromJsonElement
|
||||
import kotlinx.serialization.json.jsonObject
|
||||
import kotlinx.serialization.json.jsonPrimitive
|
||||
import kotlinx.serialization.json.long
|
||||
import kotlin.coroutines.cancellation.CancellationException
|
||||
import kotlin.test.Test
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
class ProtocolOutboundCancellationTest {
|
||||
|
||||
@Test
|
||||
fun `request times out awaiting the response and sends CancelledNotification with timeout reason`() = runTest {
|
||||
val protocol = TestProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
val inFlight = async {
|
||||
shouldThrow<McpException> {
|
||||
protocol.request<EmptyResult>(PingRequest(), RequestOptions(timeout = 5.seconds))
|
||||
}
|
||||
}
|
||||
val sent = transport.awaitRequest() // request is on the wire; peer never responds
|
||||
advanceTimeBy(6.seconds)
|
||||
runCurrent()
|
||||
|
||||
val thrown = inFlight.await()
|
||||
thrown.code shouldBe RPCError.ErrorCode.REQUEST_TIMEOUT
|
||||
thrown.data?.jsonObject?.get("timeout")?.jsonPrimitive?.long shouldBe 5.seconds.inWholeMilliseconds
|
||||
|
||||
val cancelledJson = transport.sentWithOptions
|
||||
.map { it.first }
|
||||
.filterIsInstance<JSONRPCNotification>()
|
||||
.single { it.method == Method.Defined.NotificationsCancelled.value }
|
||||
val cancelledParams = McpJson.decodeFromJsonElement<CancelledNotificationParams>(cancelledJson.params!!)
|
||||
cancelledParams.requestId shouldBe sent.id
|
||||
cancelledParams.reason.shouldNotBeNull() shouldContain "timed out"
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `outer withTimeout around request propagates the original timeout exception and notifies the peer`() = runTest {
|
||||
val protocol = TestProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
val inFlight = async {
|
||||
shouldThrow<TimeoutCancellationException> {
|
||||
withTimeout(1.seconds) {
|
||||
protocol.request<EmptyResult>(PingRequest())
|
||||
}
|
||||
}
|
||||
}
|
||||
transport.awaitRequest() // request is on the wire; peer never responds
|
||||
advanceTimeBy(2.seconds)
|
||||
runCurrent()
|
||||
inFlight.await()
|
||||
|
||||
transport.sentWithOptions.map { it.first }.filterIsInstance<JSONRPCNotification>()
|
||||
.filter { it.method == Method.Defined.NotificationsCancelled.value } shouldHaveSize 1
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `caller cancellation sends CancelledNotification and rethrows the original exception`() = runTest {
|
||||
val protocol = TestProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
val cancelled = CompletableDeferred<CancellationException>()
|
||||
val job = launch {
|
||||
try {
|
||||
protocol.request<EmptyResult>(PingRequest())
|
||||
} catch (e: CancellationException) {
|
||||
cancelled.complete(e)
|
||||
throw e
|
||||
}
|
||||
}
|
||||
transport.awaitRequest()
|
||||
job.cancelAndJoin()
|
||||
|
||||
cancelled.await()
|
||||
val wire = transport.sentWithOptions.map { it.first }.filterIsInstance<JSONRPCNotification>()
|
||||
.filter { it.method == Method.Defined.NotificationsCancelled.value }
|
||||
wire shouldHaveSize 1
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `cancelling the initialize request performs local cleanup but sends no CancelledNotification`() = runTest {
|
||||
val protocol = TestProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
val job = launch {
|
||||
protocol.request<EmptyResult>(
|
||||
InitializeRequest(
|
||||
InitializeRequestParams(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ClientCapabilities(),
|
||||
clientInfo = Implementation(name = "t", version = "1"),
|
||||
),
|
||||
),
|
||||
)
|
||||
}
|
||||
val sent = transport.awaitRequest()
|
||||
sent.method shouldBe Method.Defined.Initialize.value
|
||||
job.cancelAndJoin()
|
||||
|
||||
transport.sentWithOptions.map { it.first }.filterIsInstance<JSONRPCNotification>()
|
||||
.filter { it.method == Method.Defined.NotificationsCancelled.value } shouldHaveSize 0
|
||||
protocol.responseHandlers shouldBe emptyMap()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `late response for a cancelled request is ignored silently`() = runTest {
|
||||
val protocol = TestProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
val job = launch { protocol.request<EmptyResult>(PingRequest()) }
|
||||
val sent = transport.awaitRequest()
|
||||
job.cancelAndJoin()
|
||||
|
||||
transport.deliver(JSONRPCResponse(id = sent.id, result = EmptyResult()))
|
||||
|
||||
protocol.errors shouldHaveSize 0 // no "unknown message ID" onError
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `response for a genuinely unknown id still reports onError`() = runTest {
|
||||
val protocol = TestProtocol()
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
|
||||
transport.deliver(JSONRPCResponse(id = RequestId(999L), result = EmptyResult()))
|
||||
|
||||
protocol.errors shouldHaveSize 1
|
||||
}
|
||||
}
|
||||
|
|
@ -5,6 +5,9 @@ import io.kotest.matchers.collections.shouldContainExactly
|
|||
import io.kotest.matchers.collections.shouldHaveSize
|
||||
import io.kotest.matchers.nulls.shouldNotBeNull
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.kotest.matchers.shouldNotBe
|
||||
import io.kotest.matchers.string.shouldContain
|
||||
import io.kotest.matchers.types.shouldBeInstanceOf
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CustomRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
|
||||
|
|
@ -13,11 +16,16 @@ import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest
|
|||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCResponse
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.McpJson
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.PingRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ProgressNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ProgressNotificationParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ProgressToken
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequestParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestMeta
|
||||
import kotlinx.coroutines.async
|
||||
import kotlinx.coroutines.channels.Channel
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonObjectBuilder
|
||||
|
|
@ -30,6 +38,7 @@ import kotlinx.serialization.json.jsonPrimitive
|
|||
import kotlin.coroutines.cancellation.CancellationException
|
||||
import kotlin.test.BeforeTest
|
||||
import kotlin.test.Test
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
class ProtocolTest {
|
||||
private lateinit var protocol: TestProtocol
|
||||
|
|
@ -185,64 +194,174 @@ class ProtocolTest {
|
|||
transport.deliver(JSONRPCResponse(sent.id, EmptyResult()))
|
||||
inFlight.await()
|
||||
}
|
||||
}
|
||||
|
||||
private class TestProtocol : Protocol(null) {
|
||||
val errors = mutableListOf<Throwable>()
|
||||
@Test
|
||||
fun `request handler receives enriched extra and ambient context element`() = runTest {
|
||||
protocol.connect(transport)
|
||||
|
||||
override fun onError(error: Throwable) {
|
||||
errors.add(error)
|
||||
var seenByParameter: RequestHandlerExtra? = null
|
||||
var seenAmbient: RequestHandlerExtra? = null
|
||||
protocol.fallbackRequestHandler = { _, extra ->
|
||||
seenByParameter = extra
|
||||
seenAmbient = currentRequestHandlerExtra()
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(7L), method = "custom/echo"))
|
||||
|
||||
seenByParameter shouldNotBe null
|
||||
seenByParameter?.requestId shouldBe RequestId(7L)
|
||||
seenByParameter?.method shouldBe Method.Custom("custom/echo")
|
||||
// the SAME instance flows through both delivery paths
|
||||
seenAmbient shouldBe seenByParameter
|
||||
}
|
||||
|
||||
override fun assertCapabilityForMethod(method: Method) {
|
||||
// noop
|
||||
}
|
||||
override fun assertNotificationCapability(method: Method) {
|
||||
// noop
|
||||
}
|
||||
override fun assertRequestHandlerCapability(method: Method) {
|
||||
// noop
|
||||
}
|
||||
}
|
||||
@Test
|
||||
fun `extra method resolves defined methods to Defined entries`() = runTest {
|
||||
protocol.connect(transport)
|
||||
|
||||
private class RecordingTransport : Transport {
|
||||
private val sentMessages = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
private var onMessageCallback: (suspend (JSONRPCMessage) -> Unit)? = null
|
||||
private var onCloseCallback: (() -> Unit)? = null
|
||||
var seenMethod: Method? = null
|
||||
protocol.setRequestHandler<PingRequest>(Method.Defined.Ping) { _, extra ->
|
||||
seenMethod = extra.method
|
||||
EmptyResult()
|
||||
}
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "ping"))
|
||||
|
||||
override suspend fun start() {
|
||||
// noop
|
||||
seenMethod shouldBe Method.Defined.Ping
|
||||
}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
sentMessages.send(message)
|
||||
@Test
|
||||
fun `extra sendNotification stamps relatedRequestId`() = runTest {
|
||||
protocol.connect(transport)
|
||||
|
||||
protocol.fallbackRequestHandler = { _, extra ->
|
||||
extra.sendNotification(
|
||||
ProgressNotification(
|
||||
ProgressNotificationParams(progressToken = ProgressToken(1L), progress = 0.5),
|
||||
),
|
||||
)
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(42L), method = "custom/progressing"))
|
||||
|
||||
val sent = transport.sentWithOptions.first { it.first is JSONRPCNotification }
|
||||
sent.second?.relatedRequestId shouldBe RequestId(42L)
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
onCloseCallback?.invoke()
|
||||
@Test
|
||||
fun `extra sendRequest stamps relatedRequestId and preserves caller options`() = runTest {
|
||||
protocol.connect(transport)
|
||||
|
||||
protocol.fallbackRequestHandler = { _, extra ->
|
||||
extra.sendRequest<EmptyResult>(
|
||||
PingRequest(),
|
||||
RequestOptions(resumptionToken = "tok", onProgress = { }, timeout = 30.seconds),
|
||||
)
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
// The serial phase runs the handler inline inside deliver(), and the handler suspends
|
||||
// awaiting the nested request's response — drive delivery from a child coroutine.
|
||||
val delivery = launch {
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(42L), method = "custom/nested"))
|
||||
}
|
||||
|
||||
val outbound = transport.awaitRequest()
|
||||
val recorded = transport.sentWithOptions
|
||||
.first { (it.first as? JSONRPCRequest)?.id == outbound.id }
|
||||
.second
|
||||
val options = recorded.shouldBeInstanceOf<RequestOptions>()
|
||||
options.relatedRequestId shouldBe RequestId(42L)
|
||||
options.resumptionToken shouldBe "tok"
|
||||
options.timeout shouldBe 30.seconds
|
||||
options.onProgress shouldNotBe null
|
||||
|
||||
transport.deliver(JSONRPCResponse(id = outbound.id, result = EmptyResult()))
|
||||
delivery.join()
|
||||
}
|
||||
|
||||
override fun onClose(block: () -> Unit) {
|
||||
onCloseCallback = block
|
||||
@Test
|
||||
fun `currentRequestHandlerExtra returns null outside handlers`() = runTest {
|
||||
currentRequestHandlerExtra() shouldBe null
|
||||
}
|
||||
|
||||
override fun onError(block: (Throwable) -> Unit) {
|
||||
// noop
|
||||
@Test
|
||||
fun `connect while already connected throws IllegalStateException`() = runTest {
|
||||
protocol.connect(transport)
|
||||
shouldThrow<IllegalStateException> {
|
||||
protocol.connect(RecordingTransport())
|
||||
}
|
||||
}
|
||||
|
||||
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {
|
||||
onMessageCallback = block
|
||||
@Test
|
||||
fun `stale onClose from a previous transport does not tear down the successor connection`() = runTest {
|
||||
protocol.connect(transport)
|
||||
val staleCloseCallback = transport.closeCallback ?: error("onClose not registered")
|
||||
|
||||
protocol.close() // disconnect first transport (fires its own doClose)
|
||||
val second = RecordingTransport()
|
||||
protocol.connect(second) // reconnect
|
||||
|
||||
staleCloseCallback() // late duplicate close signal from transport #1
|
||||
|
||||
// successor connection must still be alive
|
||||
protocol.transport shouldBe second
|
||||
}
|
||||
|
||||
suspend fun awaitRequest(): JSONRPCRequest {
|
||||
val message = sentMessages.receive()
|
||||
return message as? JSONRPCRequest
|
||||
?: error("Expected JSONRPCRequest but received ${message::class.simpleName}")
|
||||
@Test
|
||||
fun `failed transport start rolls back the connection and allows reconnect`() = runTest {
|
||||
val failing = object : Transport {
|
||||
override suspend fun start(): Unit = error("boom")
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {}
|
||||
override suspend fun close() {}
|
||||
override fun onClose(block: () -> Unit) {}
|
||||
override fun onError(block: (Throwable) -> Unit) {}
|
||||
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {}
|
||||
}
|
||||
shouldThrow<IllegalStateException> { protocol.connect(failing) }
|
||||
protocol.transport shouldBe null // rolled back
|
||||
protocol.connect(transport) // reconnect succeeds, no "already connected"
|
||||
protocol.transport shouldBe transport
|
||||
}
|
||||
|
||||
suspend fun deliver(message: JSONRPCMessage) {
|
||||
val callback = onMessageCallback ?: error("onMessage callback not registered")
|
||||
callback(message)
|
||||
@Test
|
||||
fun `request cleans up its handlers when the transport send fails`() = runTest {
|
||||
// Send fails without closing the connection, so the leak is not masked by doClose().
|
||||
val failingSend = object : Transport {
|
||||
override suspend fun start() {}
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?): Unit =
|
||||
throw IllegalStateException("send failed")
|
||||
override suspend fun close() {}
|
||||
override fun onClose(block: () -> Unit) {}
|
||||
override fun onError(block: (Throwable) -> Unit) {}
|
||||
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {}
|
||||
}
|
||||
protocol.connect(failingSend)
|
||||
|
||||
shouldThrow<IllegalStateException> {
|
||||
protocol.request<EmptyResult>(PingRequest(), RequestOptions(onProgress = {}))
|
||||
}
|
||||
|
||||
protocol.responseHandlers.size shouldBe 0
|
||||
protocol.progressHandlers.size shouldBe 0
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `connect rejects non-positive maxConcurrentHandlers`() = runTest {
|
||||
val error = shouldThrow<IllegalArgumentException> {
|
||||
TestProtocol(ProtocolOptions(maxConcurrentHandlers = 0)).connect(RecordingTransport())
|
||||
}
|
||||
error.message.shouldNotBeNull() shouldContain "maxConcurrentHandlers"
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `connect rejects maxInFlightHandlers below maxConcurrentHandlers`() = runTest {
|
||||
val error = shouldThrow<IllegalArgumentException> {
|
||||
TestProtocol(ProtocolOptions(maxConcurrentHandlers = 4, maxInFlightHandlers = 0))
|
||||
.connect(RecordingTransport())
|
||||
}
|
||||
error.message.shouldNotBeNull() shouldContain "maxInFlightHandlers"
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,112 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError
|
||||
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.McpJson
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import kotlinx.coroutines.channels.Channel
|
||||
import kotlinx.serialization.json.encodeToJsonElement
|
||||
|
||||
/**
|
||||
* Shared [Protocol] test double: records [onError] calls, counts `notifications/initialized`
|
||||
* router-hook invocations, and no-ops all capability assertions.
|
||||
*
|
||||
* The concurrency gate is flipped only via [enableConcurrency]; nothing in this double flips it
|
||||
* implicitly, so tests keep today's serial semantics unless they opt in.
|
||||
*/
|
||||
internal class TestProtocol(options: ProtocolOptions? = null) : Protocol(options) {
|
||||
val errors = mutableListOf<Throwable>()
|
||||
var initializedNotificationCount = 0
|
||||
|
||||
fun enableConcurrency() {
|
||||
enableConcurrentDispatch()
|
||||
}
|
||||
|
||||
override fun onInitializedNotification() {
|
||||
initializedNotificationCount++
|
||||
// deliberately does NOT flip the gate — tests control the flip explicitly
|
||||
}
|
||||
|
||||
override fun onError(error: Throwable) {
|
||||
errors.add(error)
|
||||
}
|
||||
|
||||
override fun assertCapabilityForMethod(method: Method) {
|
||||
// noop
|
||||
}
|
||||
override fun assertNotificationCapability(method: Method) {
|
||||
// noop
|
||||
}
|
||||
public override fun assertRequestHandlerCapability(method: Method) {
|
||||
// noop
|
||||
}
|
||||
}
|
||||
|
||||
/** Shared [Transport] test double: records every sent message (with its options) and replays inbound ones. */
|
||||
internal class RecordingTransport : Transport {
|
||||
val sentMessages = Channel<JSONRPCMessage>(Channel.UNLIMITED)
|
||||
private var onMessageCallback: (suspend (JSONRPCMessage) -> Unit)? = null
|
||||
private var onCloseCallback: (() -> Unit)? = null
|
||||
|
||||
val sentWithOptions = mutableListOf<Pair<JSONRPCMessage, TransportSendOptions?>>()
|
||||
|
||||
var closeCallback: (() -> Unit)? = null
|
||||
private set
|
||||
|
||||
override suspend fun start() {
|
||||
// noop
|
||||
}
|
||||
|
||||
override suspend fun send(message: JSONRPCMessage, options: TransportSendOptions?) {
|
||||
sentWithOptions.add(message to options)
|
||||
sentMessages.send(message)
|
||||
}
|
||||
|
||||
override suspend fun close() {
|
||||
onCloseCallback?.invoke()
|
||||
}
|
||||
|
||||
override fun onClose(block: () -> Unit) {
|
||||
closeCallback = block
|
||||
onCloseCallback = block
|
||||
}
|
||||
|
||||
override fun onError(block: (Throwable) -> Unit) {
|
||||
// noop
|
||||
}
|
||||
|
||||
override fun onMessage(block: suspend (JSONRPCMessage) -> Unit) {
|
||||
onMessageCallback = block
|
||||
}
|
||||
|
||||
suspend fun awaitRequest(): JSONRPCRequest {
|
||||
while (true) {
|
||||
val message = sentMessages.receive()
|
||||
if (message is JSONRPCRequest) return message
|
||||
}
|
||||
}
|
||||
|
||||
/** Cross-thread-safe reader used by real-dispatcher tests; [sentWithOptions] is single-thread only. */
|
||||
suspend fun awaitSent(): JSONRPCMessage = sentMessages.receive()
|
||||
|
||||
suspend fun deliver(message: JSONRPCMessage) {
|
||||
val callback = onMessageCallback ?: error("onMessage callback not registered")
|
||||
callback(message)
|
||||
}
|
||||
}
|
||||
|
||||
internal fun responsesOn(transport: RecordingTransport): List<JSONRPCResponse> =
|
||||
transport.sentWithOptions.map { it.first }.filterIsInstance<JSONRPCResponse>()
|
||||
|
||||
internal fun errorsOn(transport: RecordingTransport): List<JSONRPCError> =
|
||||
transport.sentWithOptions.map { it.first }.filterIsInstance<JSONRPCError>()
|
||||
|
||||
internal fun cancelledNotification(requestId: RequestId, reason: String): JSONRPCNotification = JSONRPCNotification(
|
||||
method = Method.Defined.NotificationsCancelled.value,
|
||||
params = McpJson.encodeToJsonElement(CancelledNotificationParams(requestId = requestId, reason = reason)),
|
||||
)
|
||||
|
|
@ -0,0 +1,114 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.shared
|
||||
|
||||
import io.kotest.matchers.shouldBe
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult
|
||||
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.RequestId
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.buildJsonObject
|
||||
import kotlinx.serialization.json.int
|
||||
import kotlinx.serialization.json.jsonPrimitive
|
||||
import kotlinx.serialization.json.put
|
||||
import kotlin.test.Test
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
/**
|
||||
* Real-dispatcher counterparts to [ProtocolConcurrencyTest].
|
||||
*
|
||||
* These run under [runBlocking] (real time, real threads on [Dispatchers.Default]) rather than
|
||||
* `runTest`: a virtual-time scheduler auto-advances past the `withTimeout` guards and cannot
|
||||
* observe genuine cross-thread parallelism, so the concurrency assertions could pass vacuously.
|
||||
*/
|
||||
class ProtocolConcurrencyJvmTest {
|
||||
|
||||
@Test
|
||||
fun `fast request completes while a slow handler is parked on Dispatchers Default`() = runBlocking {
|
||||
val protocol = TestProtocol(ProtocolOptions(handlerCoroutineContext = Dispatchers.Default))
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
val slowEntered = CompletableDeferred<Unit>()
|
||||
val releaseSlow = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/slow") {
|
||||
slowEntered.complete(Unit)
|
||||
releaseSlow.await()
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow"))
|
||||
withTimeout(5.seconds) { slowEntered.await() }
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast"))
|
||||
|
||||
val fast = withTimeout(5.seconds) { transport.awaitSent() }
|
||||
(fast as JSONRPCResponse).id shouldBe RequestId(2L)
|
||||
|
||||
releaseSlow.complete(Unit)
|
||||
val slow = withTimeout(5.seconds) { transport.awaitSent() }
|
||||
(slow as JSONRPCResponse).id shouldBe RequestId(1L)
|
||||
Unit
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `dispatch stays serial before the gate flips on Dispatchers Default`() = runBlocking {
|
||||
val protocol = TestProtocol(ProtocolOptions(handlerCoroutineContext = Dispatchers.Default))
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
// gate NOT flipped
|
||||
|
||||
val slowEntered = CompletableDeferred<Unit>()
|
||||
val releaseSlow = CompletableDeferred<Unit>()
|
||||
protocol.fallbackRequestHandler = { request, _ ->
|
||||
if (request.method == "test/slow") {
|
||||
slowEntered.complete(Unit)
|
||||
releaseSlow.await()
|
||||
}
|
||||
EmptyResult()
|
||||
}
|
||||
|
||||
// A single delivering coroutine models the transport read loop: the second deliver cannot
|
||||
// start until the first returns, so a parked slow handler blocks the fast request.
|
||||
val delivery = launch {
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(1L), method = "test/slow"))
|
||||
transport.deliver(JSONRPCRequest(id = RequestId(2L), method = "test/fast"))
|
||||
}
|
||||
withTimeout(5.seconds) { slowEntered.await() }
|
||||
|
||||
// The delivering coroutine is parked inside the slow handler; fast was never delivered.
|
||||
transport.sentMessages.tryReceive().isFailure shouldBe true
|
||||
|
||||
releaseSlow.complete(Unit)
|
||||
val first = withTimeout(5.seconds) { transport.awaitSent() }
|
||||
(first as JSONRPCResponse).id shouldBe RequestId(1L)
|
||||
val second = withTimeout(5.seconds) { transport.awaitSent() }
|
||||
(second as JSONRPCResponse).id shouldBe RequestId(2L)
|
||||
delivery.join()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `suspension-free notification handlers keep arrival order on Dispatchers Default`() = runBlocking {
|
||||
val protocol = TestProtocol(ProtocolOptions(handlerCoroutineContext = Dispatchers.Default))
|
||||
val transport = RecordingTransport()
|
||||
protocol.connect(transport)
|
||||
protocol.enableConcurrency()
|
||||
|
||||
// safe: appended only from the delivering coroutine (UNDISPATCHED, suspension-free)
|
||||
val seen = mutableListOf<Int>()
|
||||
protocol.fallbackNotificationHandler = { notification ->
|
||||
seen.add((notification.params as JsonObject).getValue("i").jsonPrimitive.int)
|
||||
}
|
||||
repeat(100) { i ->
|
||||
transport.deliver(JSONRPCNotification(method = "test/event", params = buildJsonObject { put("i", i) }))
|
||||
}
|
||||
seen shouldBe (0 until 100).toList()
|
||||
}
|
||||
}
|
||||
|
|
@ -169,8 +169,8 @@ public final class io/modelcontextprotocol/kotlin/sdk/server/ServerOptions : io/
|
|||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;)V
|
||||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;Z)V
|
||||
public synthetic fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;)V
|
||||
public synthetic fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;Lkotlin/coroutines/CoroutineContext;)V
|
||||
public synthetic fun <init> (Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;ZLio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;Lkotlin/coroutines/CoroutineContext;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
|
||||
public final fun getCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/types/ServerCapabilities;
|
||||
public final fun getResourceTemplateMatcherFactory ()Lio/modelcontextprotocol/kotlin/sdk/utils/ResourceTemplateMatcherFactory;
|
||||
}
|
||||
|
|
@ -198,6 +198,7 @@ public class io/modelcontextprotocol/kotlin/sdk/server/ServerSession : io/modelc
|
|||
public fun onClose ()V
|
||||
public final fun onClose (Lkotlin/jvm/functions/Function0;)V
|
||||
public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V
|
||||
protected fun onInitializedNotification ()V
|
||||
public final fun ping (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public final fun sendElicitationComplete (Lio/modelcontextprotocol/kotlin/sdk/types/ElicitationCompleteNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
public final fun sendLoggingMessage (Lio/modelcontextprotocol/kotlin/sdk/types/LoggingMessageNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
|
||||
|
|
|
|||
|
|
@ -53,9 +53,11 @@ import io.modelcontextprotocol.kotlin.sdk.utils.ResourceTemplateMatcher
|
|||
import io.modelcontextprotocol.kotlin.sdk.utils.ResourceTemplateMatcherFactory
|
||||
import kotlinx.coroutines.CancellationException
|
||||
import kotlinx.coroutines.Deferred
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.buildJsonObject
|
||||
import kotlinx.serialization.json.put
|
||||
import kotlin.coroutines.CoroutineContext
|
||||
import kotlin.jvm.JvmOverloads
|
||||
import kotlin.time.ExperimentalTime
|
||||
|
||||
|
|
@ -69,12 +71,17 @@ private val logger = KotlinLogging.logger {}
|
|||
* Forwarded to [ProtocolOptions.enforceStrictCapabilities].
|
||||
* @property resourceTemplateMatcherFactory The factory used to create [ResourceTemplateMatcher] instances
|
||||
* for matching resource URIs against registered templates. Defaults to [PathSegmentTemplateMatcher.factory].
|
||||
* @param handlerCoroutineContext Coroutine context for inbound handlers. See [ProtocolOptions.handlerCoroutineContext].
|
||||
*/
|
||||
public class ServerOptions(
|
||||
public val capabilities: ServerCapabilities,
|
||||
enforceStrictCapabilities: Boolean = true,
|
||||
public val resourceTemplateMatcherFactory: ResourceTemplateMatcherFactory = PathSegmentTemplateMatcher.factory,
|
||||
) : ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities) {
|
||||
handlerCoroutineContext: CoroutineContext = Dispatchers.Default,
|
||||
) : ProtocolOptions(
|
||||
enforceStrictCapabilities = enforceStrictCapabilities,
|
||||
handlerCoroutineContext = handlerCoroutineContext,
|
||||
) {
|
||||
@JvmOverloads
|
||||
public constructor(
|
||||
capabilities: ServerCapabilities,
|
||||
|
|
|
|||
|
|
@ -108,10 +108,18 @@ public open class ServerSession(
|
|||
}
|
||||
}
|
||||
|
||||
override fun onInitializedNotification() {
|
||||
// Enable concurrency before the user's onInitialized callback, which runs synchronously.
|
||||
enableConcurrentDispatch()
|
||||
}
|
||||
|
||||
internal val clientConnection: ClientConnection = ClientConnectionImpl(this)
|
||||
|
||||
/**
|
||||
* Registers a callback to be invoked when the server has completed initialization.
|
||||
*
|
||||
* The callback must be synchronous and fast: it runs on the message-dispatch path for
|
||||
* `notifications/initialized`, after concurrent dispatch has been enabled for the session.
|
||||
*/
|
||||
public fun onInitialized(block: () -> Unit) {
|
||||
val old = _onInitialized
|
||||
|
|
|
|||
|
|
@ -15,10 +15,12 @@ import io.ktor.server.sse.ServerSSESession
|
|||
import io.ktor.util.collections.ConcurrentMap
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.TransportSendOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CancelledNotificationParams
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.DEFAULT_NEGOTIATED_PROTOCOL_VERSION
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCEmptyMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError
|
||||
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.McpJson
|
||||
|
|
@ -27,6 +29,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.RPCError
|
|||
import io.modelcontextprotocol.kotlin.sdk.types.RPCError.ErrorCode.REQUEST_TIMEOUT
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestId
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
import kotlinx.coroutines.NonCancellable
|
||||
import kotlinx.coroutines.awaitCancellation
|
||||
import kotlinx.coroutines.job
|
||||
|
|
@ -175,6 +178,20 @@ public class StreamableHttpServerTransport(private val configuration: Configurat
|
|||
private val requestToStreamMapping: ConcurrentMap<RequestId, String> = ConcurrentMap()
|
||||
private val requestToResponseMapping: ConcurrentMap<RequestId, JSONRPCMessage> = ConcurrentMap()
|
||||
|
||||
/**
|
||||
* JSON-response mode only: one entry per in-flight POST, completed by [send] with the collected
|
||||
* responses once every request delivered by that POST has been answered. [handlePostRequest]
|
||||
* awaits it so the HTTP exchange stays open while handlers run concurrently.
|
||||
*/
|
||||
private val pendingJsonResponses: ConcurrentMap<String, CompletableDeferred<List<JSONRPCMessage>>> = ConcurrentMap()
|
||||
|
||||
/**
|
||||
* JSON-response mode only: request ids retired by an inbound cancellation while their POST is
|
||||
* still awaiting. Treated as settled-without-response so the POST's "a response for every id"
|
||||
* completion condition is not blocked by a request that will never be answered.
|
||||
*/
|
||||
private val cancelledRequestIds: ConcurrentMap<RequestId, Unit> = ConcurrentMap()
|
||||
|
||||
private val sessionMutex = Mutex()
|
||||
private val streamMutex = Mutex()
|
||||
|
||||
|
|
@ -241,54 +258,117 @@ public class StreamableHttpServerTransport(private val configuration: Configurat
|
|||
return
|
||||
}
|
||||
|
||||
val isTerminated = message is JSONRPCResponse || message is JSONRPCError
|
||||
|
||||
// In JSON mode request-related notifications have no per-request SSE stream. Route them to the
|
||||
// standalone GET stream before the lookup below, so a notification for an already-retired
|
||||
// request is delivered there instead of becoming a spurious error.
|
||||
if (configuration.enableJsonResponse && !isTerminated) {
|
||||
val standaloneStream = streamsMapping[STANDALONE_SSE_STREAM_ID]
|
||||
standaloneStream?.let { emitOnStream(STANDALONE_SSE_STREAM_ID, it.session, message) }
|
||||
return
|
||||
}
|
||||
|
||||
val streamId = requestToStreamMapping[routingRequestId]
|
||||
?: error("No connection established for request id $routingRequestId")
|
||||
if (streamId == null) {
|
||||
// Request already retired (completed, cancelled, or the JSON POST disconnected). In JSON
|
||||
// mode a late terminal response is undeliverable and dropped. In SSE mode an unknown id
|
||||
// is a real routing error.
|
||||
if (configuration.enableJsonResponse) return
|
||||
error("No connection established for request id $routingRequestId")
|
||||
}
|
||||
val activeStream = streamsMapping[streamId]
|
||||
|
||||
if (!configuration.enableJsonResponse) {
|
||||
activeStream?.let { stream ->
|
||||
emitOnStream(streamId, stream.session, message)
|
||||
}
|
||||
}
|
||||
|
||||
val isTerminated = message is JSONRPCResponse || message is JSONRPCError
|
||||
if (!isTerminated) {
|
||||
if (configuration.enableJsonResponse) {
|
||||
// In JSON response mode there is no per-request SSE stream, so route notifications
|
||||
// that are logically associated with a request to the standalone GET SSE stream.
|
||||
val standaloneStream = streamsMapping[STANDALONE_SSE_STREAM_ID]
|
||||
standaloneStream?.let { emitOnStream(STANDALONE_SSE_STREAM_ID, it.session, message) }
|
||||
}
|
||||
return
|
||||
if (!isTerminated) return
|
||||
}
|
||||
|
||||
requestToResponseMapping[responseRequestId!!] = message
|
||||
val relatedIds = requestToStreamMapping.filterValues { it == streamId }.keys
|
||||
|
||||
if (relatedIds.any { it !in requestToResponseMapping }) return
|
||||
|
||||
streamMutex.withLock {
|
||||
if (activeStream == null) error("No connection established for request ID: $routingRequestId")
|
||||
// The request may have been retired (disconnect or cancellation cleanup) between recording
|
||||
// the response above and acquiring the lock. Reclaim the stray mapping and drop it.
|
||||
if (routingRequestId !in requestToStreamMapping) {
|
||||
requestToResponseMapping.remove(routingRequestId)
|
||||
return
|
||||
}
|
||||
// Recomputed under the lock. Handlers may respond concurrently, and the sender that
|
||||
// observes the last settled request also retires the ids, so a losing sender sees an
|
||||
// empty set. A cancelled request counts as settled even though it produces no response.
|
||||
val relatedIds = requestToStreamMapping.filterValues { it == streamId }.keys
|
||||
if (relatedIds.isEmpty() ||
|
||||
relatedIds.any { it !in requestToResponseMapping && it !in cancelledRequestIds }
|
||||
) {
|
||||
return
|
||||
}
|
||||
|
||||
if (configuration.enableJsonResponse) {
|
||||
activeStream.call.response.header(HttpHeaders.ContentType, ContentType.Application.Json.toString())
|
||||
sessionId?.let { activeStream.call.response.header(MCP_SESSION_ID_HEADER, it) }
|
||||
val responses = relatedIds.mapNotNull { requestToResponseMapping[it] }
|
||||
val payload = if (responses.size == 1) {
|
||||
responses.first()
|
||||
} else {
|
||||
responses
|
||||
// Hand the collected responses to the POST suspended in [handlePostRequest], which
|
||||
// completes the HTTP exchange. A null pendingPost means the POST was already retired
|
||||
// (disconnect or cancellation), so the responses are undeliverable and dropped quietly.
|
||||
val pendingPost = pendingJsonResponses.remove(streamId)
|
||||
if (pendingPost == null) {
|
||||
cleanUpRequests(relatedIds)
|
||||
return
|
||||
}
|
||||
activeStream.call.respond(payload)
|
||||
val responses = relatedIds.mapNotNull { requestToResponseMapping[it] }
|
||||
cleanUpRequests(relatedIds)
|
||||
pendingPost.complete(responses)
|
||||
} else {
|
||||
activeStream.session?.close()
|
||||
// A null stream means the per-request SSE connection closed before the response
|
||||
// settled, so the ids are retired rather than leaked.
|
||||
activeStream?.session?.close()
|
||||
cleanUpRequests(relatedIds)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up
|
||||
relatedIds.forEach { requestId ->
|
||||
requestToResponseMapping.remove(requestId)
|
||||
requestToStreamMapping.remove(requestId)
|
||||
private fun cleanUpRequests(ids: Collection<RequestId>) {
|
||||
ids.forEach { requestId ->
|
||||
requestToResponseMapping.remove(requestId)
|
||||
requestToStreamMapping.remove(requestId)
|
||||
cancelledRequestIds.remove(requestId)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Scans a delivered POST body for `notifications/cancelled` and retires each referenced request
|
||||
* that is awaiting on a JSON-mode POST. No-op unless [Configuration.enableJsonResponse].
|
||||
*/
|
||||
private suspend fun retireCancelledJsonRequests(messages: List<JSONRPCMessage>) {
|
||||
if (!configuration.enableJsonResponse) return
|
||||
val cancelledIds = messages.asSequence()
|
||||
.filterIsInstance<JSONRPCNotification>()
|
||||
.filter { it.method == Method.Defined.NotificationsCancelled.value }
|
||||
.mapNotNull { notification ->
|
||||
val params = notification.params ?: return@mapNotNull null
|
||||
runCatching {
|
||||
McpJson.decodeFromJsonElement<CancelledNotificationParams>(params).requestId
|
||||
}.getOrNull()
|
||||
}
|
||||
.toList()
|
||||
cancelledIds.forEach { retireCancelledJsonRequest(it) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Marks a request cancelled by the peer as settled without a response and completes its POST once
|
||||
* every request that POST delivered is settled (answered or cancelled). Without this, a cancelled
|
||||
* request leaves the POST awaiting a response that, by design, is never sent.
|
||||
*/
|
||||
private suspend fun retireCancelledJsonRequest(requestId: RequestId) {
|
||||
streamMutex.withLock {
|
||||
val streamId = requestToStreamMapping[requestId] ?: return
|
||||
cancelledRequestIds[requestId] = Unit
|
||||
val relatedIds = requestToStreamMapping.filterValues { it == streamId }.keys
|
||||
val allSettled = relatedIds.all { it in requestToResponseMapping || it in cancelledRequestIds }
|
||||
if (!allSettled) return
|
||||
val pendingPost = pendingJsonResponses.remove(streamId) ?: return
|
||||
val responses = relatedIds.mapNotNull { requestToResponseMapping[it] }
|
||||
cleanUpRequests(relatedIds)
|
||||
pendingPost.complete(responses)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -304,6 +384,10 @@ public class StreamableHttpServerTransport(private val configuration: Configurat
|
|||
streamsMapping.clear()
|
||||
requestToStreamMapping.clear()
|
||||
requestToResponseMapping.clear()
|
||||
cancelledRequestIds.clear()
|
||||
// Release any POST calls suspended awaiting JSON responses.
|
||||
pendingJsonResponses.values.forEach { it.cancel() }
|
||||
pendingJsonResponses.clear()
|
||||
invokeOnCloseCallback()
|
||||
}
|
||||
}
|
||||
|
|
@ -339,7 +423,12 @@ public class StreamableHttpServerTransport(private val configuration: Configurat
|
|||
}
|
||||
|
||||
/**
|
||||
* Handles POST requests containing JSON-RPC messages
|
||||
* Handles POST requests containing JSON-RPC messages.
|
||||
*
|
||||
* In JSON-response mode this suspends until a response has been produced for every request in
|
||||
* the POST body, then completes the HTTP exchange with the collected JSON (a single object for
|
||||
* one request, an array for a batch). Bodies without requests are acknowledged with 202
|
||||
* immediately. In SSE mode the responses are emitted onto the per-request SSE stream instead.
|
||||
*/
|
||||
public suspend fun handlePostRequest(session: ServerSSESession?, call: ApplicationCall) {
|
||||
try {
|
||||
|
|
@ -407,6 +496,9 @@ public class StreamableHttpServerTransport(private val configuration: Configurat
|
|||
if (!hasRequest) {
|
||||
call.respondNullable(status = HttpStatusCode.Accepted, message = null)
|
||||
messages.forEach { message -> _onMessage(message) }
|
||||
// A cancellation may target a request still awaiting delivery on a JSON-mode POST;
|
||||
// retire it so that POST completes instead of waiting for a response that never comes.
|
||||
retireCancelledJsonRequests(messages)
|
||||
return
|
||||
}
|
||||
|
||||
|
|
@ -429,13 +521,50 @@ public class StreamableHttpServerTransport(private val configuration: Configurat
|
|||
maybeSendPrimingEvent(streamId, session, clientProtocolVersion)
|
||||
}
|
||||
|
||||
val pendingPost = if (configuration.enableJsonResponse) {
|
||||
CompletableDeferred<List<JSONRPCMessage>>()
|
||||
} else {
|
||||
null
|
||||
}
|
||||
|
||||
streamMutex.withLock {
|
||||
streamsMapping[streamId] = SessionContext(session, call)
|
||||
messages.filterIsInstance<JSONRPCRequest>().forEach { requestToStreamMapping[it.id] = streamId }
|
||||
pendingPost?.let { pendingJsonResponses[streamId] = it }
|
||||
}
|
||||
call.coroutineContext.job.invokeOnCompletion { streamsMapping.remove(streamId) }
|
||||
|
||||
messages.forEach { message -> _onMessage(message) }
|
||||
if (pendingPost == null) {
|
||||
// SSE mode: responses are emitted onto this call's SSE stream by send().
|
||||
messages.forEach { message -> _onMessage(message) }
|
||||
return
|
||||
}
|
||||
|
||||
// JSON-response mode. Handlers may run concurrently, so returning from delivery no longer
|
||||
// implies the responses were produced. Suspend until send() has collected a response for
|
||||
// every request in this POST. A client disconnect or transport close cancels this call.
|
||||
try {
|
||||
messages.forEach { message -> _onMessage(message) }
|
||||
retireCancelledJsonRequests(messages)
|
||||
val responses = pendingPost.await()
|
||||
sessionId?.let { call.response.header(MCP_SESSION_ID_HEADER, it) }
|
||||
if (responses.isEmpty()) {
|
||||
// Every request in this POST was cancelled before producing a response; there is
|
||||
// nothing to return, so acknowledge with 202 rather than an empty JSON body.
|
||||
call.respondNullable(status = HttpStatusCode.Accepted, message = null)
|
||||
} else {
|
||||
call.response.header(HttpHeaders.ContentType, ContentType.Application.Json.toString())
|
||||
val payload = if (responses.size == 1) responses.first() else responses
|
||||
call.respond(payload)
|
||||
}
|
||||
} finally {
|
||||
// On cancellation or failure, retire this POST's ids so a late response is dropped
|
||||
// quietly at the send() boundary instead of leaking mappings. Idempotent no-ops on
|
||||
// the success path (send() or retirement already cleaned up).
|
||||
pendingJsonResponses.remove(streamId)
|
||||
pendingPost.cancel()
|
||||
cleanUpRequests(messages.filterIsInstance<JSONRPCRequest>().map { it.id })
|
||||
}
|
||||
} catch (e: CancellationException) {
|
||||
throw e
|
||||
} catch (e: Exception) {
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue