From 7cf7eb78acac88c06abe0ea890d6fe075c1ab5a8 Mon Sep 17 00:00:00 2001 From: taekop Date: Wed, 23 Sep 2026 10:51:18 +0900 Subject: [PATCH 1/3] fix(server): run onInitialized callbacks registered after initialization (#920) ServerSession.onInitialized composed callbacks into a plain var, so a callback registered after notifications/initialized was never invoked, and concurrent registration could drop callbacks. Keep pending callbacks in an AtomicRef. The initialized handler swaps it to null and runs the list in registration order; registration after that runs the callback immediately. Each callback now runs at most once. Fixes #920 --- .../sdk/server/ServerSessionInitializeTest.kt | 47 +++++++++++++++++++ .../kotlin/sdk/server/ServerSession.kt | 22 ++++++--- 2 files changed, 62 insertions(+), 7 deletions(-) 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..26ea70199 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 @@ -30,6 +31,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 +163,49 @@ class ServerSessionInitializeTest { successes shouldHaveSize 1 errors.map { it.error.code } shouldBe List(n - 1) { RPCError.ErrorCode.INVALID_REQUEST } } + + private suspend fun completeHandshake(session: ServerSession) { + 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() + clientTransport.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 every onInitialized callback registered concurrently`() = runTest { + val session = createSession() + val counters = List(200) { AtomicInteger() } + + withContext(Dispatchers.Default) { + counters.map { counter -> + launch { session.onInitialized { counter.incrementAndGet() } } + }.joinAll() + } + completeHandshake(session) + + counters.map { it.get() } shouldBe List(counters.size) { 1 } + } } 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..94b3f71a7 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,6 +31,9 @@ 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.uuid.ExperimentalUuidApi @@ -64,7 +67,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 +102,7 @@ public open class ServerSession( handleInitialize(request) } setNotificationHandler(Defined.NotificationsInitialized) { - _onInitialized() + pendingInitializedCallbacks.getAndSet(null)?.forEach { callback -> callback() } CompletableDeferred(Unit) } @@ -120,13 +128,13 @@ 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. */ public fun onInitialized(block: () -> Unit) { - val old = _onInitialized - _onInitialized = { - old() - block() - } + if (pendingInitializedCallbacks.getAndUpdate { it?.adding(block) } == null) block() } /** From 51124e66e773ccd1ab5105d80b9e0973177150cb Mon Sep 17 00:00:00 2001 From: devcrocod Date: Wed, 14 Oct 2026 18:36:26 +0200 Subject: [PATCH 2/3] Handle exceptions in `onInitialized` callbacks and enhance `ServerSession` tests. --- .../sdk/server/ServerSessionInitializeTest.kt | 54 ++++++++++++++++--- .../kotlin/sdk/server/ServerSession.kt | 14 ++++- 2 files changed, 60 insertions(+), 8 deletions(-) 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 26ea70199..4767e3992 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 @@ -22,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 @@ -164,7 +165,7 @@ class ServerSessionInitializeTest { errors.map { it.error.code } shouldBe List(n - 1) { RPCError.ErrorCode.INVALID_REQUEST } } - private suspend fun completeHandshake(session: ServerSession) { + private suspend fun connectAndInitialize(session: ServerSession): InMemoryTransport { val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair() val responseDone = CompletableDeferred() clientTransport.onMessage { message -> @@ -174,7 +175,11 @@ class ServerSessionInitializeTest { session.connect(serverTransport) clientTransport.send(createInitializeRequest("test-client").toJSON()) responseDone.await() - clientTransport.send(InitializedNotification().toJSON()) + return clientTransport + } + + private suspend fun completeHandshake(session: ServerSession) { + connectAndInitialize(session).send(InitializedNotification().toJSON()) } @Test @@ -195,17 +200,52 @@ class ServerSessionInitializeTest { } @Test - fun `should run every onInitialized callback registered concurrently`() = runTest { + 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) { - counters.map { counter -> - launch { session.onInitialized { counter.incrementAndGet() } } - }.joinAll() + 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() } - completeHandshake(session) + // 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 a late onInitialized callback failure to onError`() = 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) + } + } + completeHandshake(session) + + session.onInitialized { error("late failure") } + + errors.map { it.message } shouldBe listOf("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 94b3f71a7..af78f6ecd 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 @@ -36,6 +36,7 @@ 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 @@ -132,9 +133,20 @@ public open class ServerSession( * 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. */ public fun onInitialized(block: () -> Unit) { - if (pendingInitializedCallbacks.getAndUpdate { it?.adding(block) } == null) block() + if (pendingInitializedCallbacks.getAndUpdate { it?.adding(block) } == null) { + try { + block() + } catch (e: CancellationException) { + throw e + } catch (cause: Throwable) { + logger.error(cause) { "Error in onInitialized callback" } + onError(cause) + } + } } /** From d59c4205572c891b748735f8642abcbc9caa90b2 Mon Sep 17 00:00:00 2001 From: devcrocod Date: Thu, 15 Oct 2026 17:50:37 +0200 Subject: [PATCH 3/3] Handle exceptions in `onInitialized` to avoid skipping other callbacks; improve related test coverage. --- .../kotlin/sdk/server/ServerSessionInitializeTest.kt | 10 +++++++--- .../kotlin/sdk/server/ServerSession.kt | 8 ++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) 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 4767e3992..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 @@ -231,7 +231,7 @@ class ServerSessionInitializeTest { } @Test - fun `should report a late onInitialized callback failure to onError`() = runTest { + 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"), @@ -242,10 +242,14 @@ class ServerSessionInitializeTest { errors.add(error) } } - completeHandshake(session) + val calls = CopyOnWriteArrayList() + session.onInitialized { error("early failure") } + session.onInitialized { calls.add("early") } + completeHandshake(session) session.onInitialized { error("late failure") } - errors.map { it.message } shouldBe listOf("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 af78f6ecd..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 @@ -134,10 +134,11 @@ public open class ServerSession( * 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. + * 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) { - if (pendingInitializedCallbacks.getAndUpdate { it?.adding(block) } == null) { + val callback = { try { block() } catch (e: CancellationException) { @@ -147,6 +148,9 @@ public open class ServerSession( onError(cause) } } + if (pendingInitializedCallbacks.getAndUpdate { it?.adding(callback) } == null) { + callback() + } } /**