fix: reject duplicate initialize requests in ServerSession (#624)
`ServerSession.handleInitialize()` now rejects duplicate `initialize` requests with `INVALID_REQUEST` (-32600) error. `clientCapabilities` and `clientVersion` changed from `var` to `AtomicRef`-backed `val`, with atomic `compareAndSet(null, ...)` serving as both the guard and the write ## How Has This Been Tested? second `initialize` returns error, original capabilities preserved ## Breaking Changes none ## Types of changes - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] 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 ## Additional context <!-- Add any other context, implementation notes, or design decisions -->
This commit is contained in:
parent
e7687c7939
commit
9c0aa45e2a
3 changed files with 172 additions and 10 deletions
|
|
@ -417,13 +417,14 @@ class ClientTest {
|
||||||
clientInfo = Implementation(name = "test client without capability", version = "1.0"),
|
clientInfo = Implementation(name = "test client without capability", version = "1.0"),
|
||||||
options = ClientOptions(
|
options = ClientOptions(
|
||||||
capabilities = ClientCapabilities(),
|
capabilities = ClientCapabilities(),
|
||||||
// enforceStrictCapabilities = true // TODO()
|
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
clientWithoutCapability.connect(clientTransport)
|
val (clientTransport2, serverTransport2) = InMemoryTransport.createLinkedPair()
|
||||||
// Using the same transport pair might not be realistic - in a real scenario you'd create another pair.
|
listOf(
|
||||||
// Adjust if necessary.
|
launch { clientWithoutCapability.connect(clientTransport2) },
|
||||||
|
launch { server.createSession(serverTransport2) },
|
||||||
|
).joinAll()
|
||||||
|
|
||||||
// This should fail
|
// This should fail
|
||||||
val ex = assertFailsWith<IllegalStateException> {
|
val ex = assertFailsWith<IllegalStateException> {
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,152 @@
|
||||||
|
package io.modelcontextprotocol.kotlin.sdk.server
|
||||||
|
|
||||||
|
import io.modelcontextprotocol.kotlin.sdk.shared.InMemoryTransport
|
||||||
|
import io.modelcontextprotocol.kotlin.sdk.types.ClientCapabilities
|
||||||
|
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.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.RPCError
|
||||||
|
import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities
|
||||||
|
import io.modelcontextprotocol.kotlin.sdk.types.toJSON
|
||||||
|
import kotlinx.coroutines.CompletableDeferred
|
||||||
|
import kotlinx.coroutines.Dispatchers
|
||||||
|
import kotlinx.coroutines.joinAll
|
||||||
|
import kotlinx.coroutines.launch
|
||||||
|
import kotlinx.coroutines.test.runTest
|
||||||
|
import kotlinx.coroutines.withContext
|
||||||
|
import org.junit.jupiter.api.Test
|
||||||
|
import java.util.concurrent.CopyOnWriteArrayList
|
||||||
|
import kotlin.test.assertEquals
|
||||||
|
import kotlin.test.assertNotNull
|
||||||
|
import kotlin.test.assertNull
|
||||||
|
import kotlin.test.assertTrue
|
||||||
|
|
||||||
|
class ServerSessionInitializeTest {
|
||||||
|
|
||||||
|
private fun createSession(): ServerSession = ServerSession(
|
||||||
|
serverInfo = Implementation(name = "test-server", version = "1.0"),
|
||||||
|
options = ServerOptions(capabilities = ServerCapabilities()),
|
||||||
|
instructions = null,
|
||||||
|
)
|
||||||
|
|
||||||
|
private fun createInitializeRequest(clientName: String = "test-client"): InitializeRequest = InitializeRequest(
|
||||||
|
InitializeRequestParams(
|
||||||
|
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||||
|
capabilities = ClientCapabilities(),
|
||||||
|
clientInfo = Implementation(name = clientName, version = "1.0"),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `should handle first initialize request successfully`() = runTest {
|
||||||
|
val session = createSession()
|
||||||
|
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||||
|
|
||||||
|
assertNull(session.clientCapabilities)
|
||||||
|
assertNull(session.clientVersion)
|
||||||
|
|
||||||
|
val responseDone = CompletableDeferred<JSONRPCResponse>()
|
||||||
|
clientTransport.onMessage { message ->
|
||||||
|
if (message is JSONRPCResponse) {
|
||||||
|
responseDone.complete(message)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
session.connect(serverTransport)
|
||||||
|
clientTransport.send(createInitializeRequest().toJSON())
|
||||||
|
|
||||||
|
val response = responseDone.await()
|
||||||
|
assertNotNull(response.result)
|
||||||
|
assertNotNull(session.clientCapabilities)
|
||||||
|
assertEquals("test-client", session.clientVersion?.name)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `should reject duplicate initialize request`() = runTest {
|
||||||
|
val session = createSession()
|
||||||
|
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||||
|
|
||||||
|
val responses = CopyOnWriteArrayList<JSONRPCMessage>()
|
||||||
|
val secondResponseDone = CompletableDeferred<Unit>()
|
||||||
|
|
||||||
|
clientTransport.onMessage { message ->
|
||||||
|
when (message) {
|
||||||
|
is JSONRPCResponse, is JSONRPCError -> {
|
||||||
|
responses.add(message)
|
||||||
|
if (responses.size == 2) secondResponseDone.complete(Unit)
|
||||||
|
}
|
||||||
|
|
||||||
|
else -> {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
session.connect(serverTransport)
|
||||||
|
|
||||||
|
// First initialize should succeed
|
||||||
|
clientTransport.send(createInitializeRequest(clientName = "first-client").toJSON())
|
||||||
|
|
||||||
|
// Second initialize should be rejected
|
||||||
|
clientTransport.send(createInitializeRequest(clientName = "second-client").toJSON())
|
||||||
|
|
||||||
|
secondResponseDone.await()
|
||||||
|
|
||||||
|
assertEquals(2, responses.size)
|
||||||
|
assertTrue(responses[0] is JSONRPCResponse, "First response should be success")
|
||||||
|
assertTrue(responses[1] is JSONRPCError, "Second response should be error")
|
||||||
|
assertEquals(RPCError.ErrorCode.INVALID_REQUEST, (responses[1] as JSONRPCError).error.code)
|
||||||
|
|
||||||
|
// Capabilities still reflect the first client, not overwritten
|
||||||
|
assertEquals("first-client", session.clientVersion?.name)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `should reject concurrent initialize requests - only first succeeds`() = runTest {
|
||||||
|
val session = createSession()
|
||||||
|
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||||
|
|
||||||
|
val n = 10
|
||||||
|
val allResponsesDone = CompletableDeferred<Unit>()
|
||||||
|
val successes = CopyOnWriteArrayList<JSONRPCResponse>()
|
||||||
|
val errors = CopyOnWriteArrayList<JSONRPCError>()
|
||||||
|
|
||||||
|
clientTransport.onMessage { message ->
|
||||||
|
when (message) {
|
||||||
|
is JSONRPCResponse -> successes.add(message)
|
||||||
|
is JSONRPCError -> errors.add(message)
|
||||||
|
else -> {}
|
||||||
|
}
|
||||||
|
if (successes.size + errors.size == n) {
|
||||||
|
allResponsesDone.complete(Unit)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
session.connect(serverTransport)
|
||||||
|
|
||||||
|
// Use Dispatchers.Default for true parallelism on JVM
|
||||||
|
withContext(Dispatchers.Default) {
|
||||||
|
val barrier = CompletableDeferred<Unit>()
|
||||||
|
val jobs = (1..n).map { i ->
|
||||||
|
launch {
|
||||||
|
barrier.await()
|
||||||
|
clientTransport.send(
|
||||||
|
createInitializeRequest(clientName = "client-$i").toJSON(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier.complete(Unit)
|
||||||
|
jobs.joinAll()
|
||||||
|
}
|
||||||
|
|
||||||
|
allResponsesDone.await()
|
||||||
|
|
||||||
|
assertEquals(1, successes.size, "Exactly one initialize should succeed")
|
||||||
|
assertEquals(n - 1, errors.size, "All other initializes should be rejected")
|
||||||
|
errors.forEach { error ->
|
||||||
|
assertEquals(RPCError.ErrorCode.INVALID_REQUEST, error.error.code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -20,8 +20,10 @@ import io.modelcontextprotocol.kotlin.sdk.types.ListRootsRequest
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.ListRootsResult
|
import io.modelcontextprotocol.kotlin.sdk.types.ListRootsResult
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.LoggingLevel
|
import io.modelcontextprotocol.kotlin.sdk.types.LoggingLevel
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.LoggingMessageNotification
|
import io.modelcontextprotocol.kotlin.sdk.types.LoggingMessageNotification
|
||||||
|
import io.modelcontextprotocol.kotlin.sdk.types.McpException
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.Method
|
import io.modelcontextprotocol.kotlin.sdk.types.Method
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.Method.Defined
|
import io.modelcontextprotocol.kotlin.sdk.types.Method.Defined
|
||||||
|
import io.modelcontextprotocol.kotlin.sdk.types.RPCError
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.RequestMeta
|
import io.modelcontextprotocol.kotlin.sdk.types.RequestMeta
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification
|
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification
|
||||||
import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS
|
import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS
|
||||||
|
|
@ -52,17 +54,18 @@ public open class ServerSession(
|
||||||
|
|
||||||
private var _onClose: () -> Unit = {}
|
private var _onClose: () -> Unit = {}
|
||||||
|
|
||||||
|
private val _clientCapabilities: AtomicRef<ClientCapabilities?> = atomic(null)
|
||||||
|
private val _clientVersion: AtomicRef<Implementation?> = atomic(null)
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* The client's reported capabilities after initialization.
|
* The client's reported capabilities after initialization.
|
||||||
*/
|
*/
|
||||||
public var clientCapabilities: ClientCapabilities? = null
|
public val clientCapabilities: ClientCapabilities? get() = _clientCapabilities.value
|
||||||
private set
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* The client's version information after initialization.
|
* The client's version information after initialization.
|
||||||
*/
|
*/
|
||||||
public var clientVersion: Implementation? = null
|
public val clientVersion: Implementation? get() = _clientVersion.value
|
||||||
private set
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* The capabilities supported by the server, related to the session.
|
* The capabilities supported by the server, related to the session.
|
||||||
|
|
@ -286,9 +289,15 @@ public open class ServerSession(
|
||||||
}
|
}
|
||||||
|
|
||||||
private fun handleInitialize(request: InitializeRequest): InitializeResult {
|
private fun handleInitialize(request: InitializeRequest): InitializeResult {
|
||||||
|
if (!_clientCapabilities.compareAndSet(null, request.params.capabilities)) {
|
||||||
|
throw McpException(
|
||||||
|
code = RPCError.ErrorCode.INVALID_REQUEST,
|
||||||
|
message = "Server already initialized",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
logger.debug { "Handling initialization request from client" }
|
logger.debug { "Handling initialization request from client" }
|
||||||
clientCapabilities = request.params.capabilities
|
_clientVersion.value = request.params.clientInfo
|
||||||
clientVersion = request.params.clientInfo
|
|
||||||
|
|
||||||
val requestedVersion = request.params.protocolVersion
|
val requestedVersion = request.params.protocolVersion
|
||||||
val protocolVersion = if (SUPPORTED_PROTOCOL_VERSIONS.contains(requestedVersion)) {
|
val protocolVersion = if (SUPPORTED_PROTOCOL_VERSIONS.contains(requestedVersion)) {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue