diff --git a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt index 2683b7527..5216adcbb 100644 --- a/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt +++ b/integration-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSessionInitializeTest.kt @@ -10,6 +10,7 @@ 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.InitializedNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCError import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCRequest @@ -21,6 +22,7 @@ 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.Job import kotlinx.coroutines.joinAll import kotlinx.coroutines.launch import kotlinx.coroutines.test.runTest @@ -30,6 +32,7 @@ import kotlinx.serialization.json.put import kotlinx.serialization.json.putJsonObject import org.junit.jupiter.api.Test import java.util.concurrent.CopyOnWriteArrayList +import java.util.concurrent.atomic.AtomicInteger class ServerSessionInitializeTest { @@ -161,4 +164,92 @@ class ServerSessionInitializeTest { successes shouldHaveSize 1 errors.map { it.error.code } shouldBe List(n - 1) { RPCError.ErrorCode.INVALID_REQUEST } } + + private suspend fun connectAndInitialize(session: ServerSession): InMemoryTransport { + val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() + val responseDone = CompletableDeferred() + clientTransport.onMessage { message -> + if (message is JSONRPCResponse) responseDone.complete(message) + } + + session.connect(serverTransport) + clientTransport.send(createInitializeRequest("test-client").toJSON()) + responseDone.await() + return clientTransport + } + + private suspend fun completeHandshake(session: ServerSession) { + connectAndInitialize(session).send(InitializedNotification().toJSON()) + } + + @Test + fun `should run onInitialized callback registered after initialization`() = runTest { + val session = createSession() + val calls = CopyOnWriteArrayList() + val initialized = CompletableDeferred() + + session.onInitialized { calls.add("early") } + session.onInitialized { initialized.complete(Unit) } + + completeHandshake(session) + initialized.await() + + session.onInitialized { calls.add("late") } + + calls shouldBe listOf("early", "late") + } + + @Test + fun `should run each onInitialized callback exactly once when registration races initialization`() = runTest { + val session = createSession() + val clientTransport = connectAndInitialize(session) + val counters = List(200) { AtomicInteger() } + + withContext(Dispatchers.Default) { + val barrier = CompletableDeferred() + val jobs = mutableListOf() + counters.forEachIndexed { index, counter -> + // Mid-batch, so registrations land before, during and after the notification. + if (index == counters.size / 2) { + jobs += launch { + barrier.await() + clientTransport.send(InitializedNotification().toJSON()) + } + } + jobs += launch { + barrier.await() + session.onInitialized { counter.incrementAndGet() } + } + } + barrier.complete(Unit) + jobs.joinAll() + } + // A duplicate notification must not run any callback again. + clientTransport.send(InitializedNotification().toJSON()) + + counters.map { it.get() } shouldBe List(counters.size) { 1 } + } + + @Test + fun `should report onInitialized callback failures to onError without skipping other callbacks`() = runTest { + val errors = CopyOnWriteArrayList() + val session = object : ServerSession( + serverInfo = Implementation(name = "test-server", version = "1.0"), + options = ServerOptions(capabilities = ServerCapabilities()), + instructions = null, + ) { + override fun onError(error: Throwable) { + errors.add(error) + } + } + val calls = CopyOnWriteArrayList() + session.onInitialized { error("early failure") } + session.onInitialized { calls.add("early") } + + completeHandshake(session) + session.onInitialized { error("late failure") } + + calls shouldBe listOf("early") + errors.map { it.message } shouldBe listOf("early failure", "late failure") + } } diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt index f20d3e122..f8a38262e 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerSession.kt @@ -31,8 +31,12 @@ import io.modelcontextprotocol.kotlin.sdk.types.SUPPORTED_PROTOCOL_VERSIONS import io.modelcontextprotocol.kotlin.sdk.types.SetLevelRequest import kotlinx.atomicfu.AtomicRef import kotlinx.atomicfu.atomic +import kotlinx.atomicfu.getAndUpdate +import kotlinx.collections.immutable.PersistentList +import kotlinx.collections.immutable.persistentListOf import kotlinx.coroutines.CompletableDeferred import kotlinx.serialization.json.JsonObject +import kotlin.coroutines.cancellation.CancellationException import kotlin.uuid.ExperimentalUuidApi import kotlin.uuid.Uuid @@ -64,7 +68,12 @@ public open class ServerSession( @OptIn(ExperimentalUuidApi::class) public val sessionId: String = Uuid.random().toString() - private var _onInitialized: (() -> Unit) = {} + /** + * Callbacks waiting for `notifications/initialized`, or `null` once it has been received. + * A callback registered after that point runs immediately instead of being stored. + */ + private val pendingInitializedCallbacks: AtomicRef Unit>?> = + atomic(persistentListOf()) private var _onClose: () -> Unit = {} @@ -94,7 +103,7 @@ public open class ServerSession( handleInitialize(request) } setNotificationHandler(Defined.NotificationsInitialized) { - _onInitialized() + pendingInitializedCallbacks.getAndSet(null)?.forEach { callback -> callback() } CompletableDeferred(Unit) } @@ -120,12 +129,27 @@ public open class ServerSession( * * 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. + * + * If initialization has already completed, [block] runs immediately on the calling thread. + * Otherwise it runs when `notifications/initialized` is received, after the callbacks + * registered before it. Each callback runs at most once. + * + * An exception thrown by [block] is reported to [onError] instead of being propagated, + * so it does not prevent other callbacks from running. */ public fun onInitialized(block: () -> Unit) { - val old = _onInitialized - _onInitialized = { - old() - block() + val callback = { + try { + block() + } catch (e: CancellationException) { + throw e + } catch (cause: Throwable) { + logger.error(cause) { "Error in onInitialized callback" } + onError(cause) + } + } + if (pendingInitializedCallbacks.getAndUpdate { it?.adding(callback) } == null) { + callback() } }