fix(server): run onInitialized callbacks registered after initialization

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<JSONRPCResponse>()
+        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<String>()
+        val initialized = CompletableDeferred<Unit>()
+
+        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<Unit>()
+            val jobs = mutableListOf<Job>()
+            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<Throwable>()
+        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<String>()
+        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")
+    }
 }

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<PersistentList<() -> Unit>?> =
+        atomic(persistentListOf())
 
     private var _onClose: () -> Unit = {}
 
@@ -94,7 +103,7 @@ public open class ServerSession(
             handleInitialize(request)
         }
         setNotificationHandler<InitializedNotification>(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()
         }
     }