Skip to content

Commit 272c2e3

Browse files
tiginamariatiginamaria
committed
Support notifications for server
1 parent 6be350d commit 272c2e3

File tree

12 files changed

+1112
-126
lines changed

12 files changed

+1112
-126
lines changed

‎kotlin-sdk-server/api/kotlin-sdk-server.api‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,8 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server {
8282
public final fun onConnect (Lkotlin/jvm/functions/Function0;)V
8383
public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V
8484
public final fun ping (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
85+
public final fun removeNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V
86+
public final fun removeNotificationHandler (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V
8587
public final fun removePrompt (Ljava/lang/String;)Z
8688
public final fun removePrompts (Ljava/util/List;)I
8789
public final fun removeResource (Ljava/lang/String;)Z
@@ -93,6 +95,8 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server {
9395
public final fun sendResourceListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
9496
public final fun sendResourceUpdated (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/ResourceUpdatedNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
9597
public final fun sendToolListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
98+
public final fun setNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;Lkotlin/jvm/functions/Function1;)V
99+
public final fun setNotificationHandler (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/Method;Lkotlin/jvm/functions/Function1;)V
96100
}
97101

98102
public final class io/modelcontextprotocol/kotlin/sdk/server/ServerOptions : io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions {

‎kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureNotificationService.kt‎

Lines changed: 329 additions & 67 deletions
Large diffs are not rendered by default.

‎kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureRegistry.kt‎

Lines changed: 20 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import kotlinx.collections.immutable.toPersistentSet
1313
* A listener interface for receiving notifications about feature changes in registry.
1414
*/
1515
internal interface FeatureListener {
16-
fun onListChanged()
1716
fun onFeatureUpdated(featureKey: String)
1817
}
1918

@@ -61,10 +60,9 @@ internal class FeatureRegistry<T : Feature>(private val featureType: String) {
6160
logger.info { "Adding $featureType: \"${feature.key}\"" }
6261
val oldMap = registry.getAndUpdate { current -> current.put(feature.key, feature) }
6362
val oldFeature = oldMap[feature.key]
64-
logger.info { "Added $featureType: \"${feature.key}\"" }
6563

64+
logger.info { "Added $featureType: \"${feature.key}\"" }
6665
notifyFeatureUpdated(oldFeature, feature)
67-
notifyListChanged()
6866
}
6967

7068
/**
@@ -76,13 +74,12 @@ internal class FeatureRegistry<T : Feature>(private val featureType: String) {
7674
internal fun addAll(features: List<T>) {
7775
logger.info { "Adding ${featureType}s: ${features.size}" }
7876
val oldMap = registry.getAndUpdate { current -> current.putAll(features.associateBy { it.key }) }
77+
78+
logger.info { "Added ${featureType}s: ${features.size}" }
7979
for (feature in features) {
8080
val oldFeature = oldMap[feature.key]
8181
notifyFeatureUpdated(oldFeature, feature)
8282
}
83-
logger.info { "Added ${featureType}s: ${features.size}" }
84-
85-
notifyListChanged()
8683
}
8784

8885
/**
@@ -97,16 +94,13 @@ internal class FeatureRegistry<T : Feature>(private val featureType: String) {
9794

9895
val removedFeature = oldMap[key]
9996
val removed = removedFeature != null
100-
logger.info {
101-
if (removed) {
102-
"Removed $featureType: \"$key\""
103-
} else {
104-
"$featureType not found: \"$key\""
105-
}
106-
}
10797

108-
notifyFeatureUpdated(removedFeature, null)
109-
notifyListChanged()
98+
if (removed) {
99+
logger.info { "Removed $featureType: \"$key\"" }
100+
notifyFeatureUpdated(removedFeature, null)
101+
} else {
102+
logger.info { "$featureType not found: \"$key\"" }
103+
}
110104

111105
return removed
112106
}
@@ -123,18 +117,16 @@ internal class FeatureRegistry<T : Feature>(private val featureType: String) {
123117

124118
val removedFeatures = keys.mapNotNull { oldMap[it] }
125119
val removedCount = removedFeatures.size
120+
121+
if (removedCount > 0) {
122+
logger.info { "Removed ${featureType}s: $removedCount" }
123+
} else {
124+
logger.info { "No $featureType were removed" }
125+
}
126126
removedFeatures.forEach {
127127
notifyFeatureUpdated(it, null)
128128
}
129129

130-
logger.info {
131-
if (removedCount > 0) {
132-
"Removed ${featureType}s: $removedCount"
133-
} else {
134-
"No $featureType were removed"
135-
}
136-
}
137-
138130
return removedCount
139131
}
140132

@@ -156,15 +148,13 @@ internal class FeatureRegistry<T : Feature>(private val featureType: String) {
156148
return feature
157149
}
158150

159-
private fun notifyListChanged() {
160-
logger.info { "Notifying listeners of list change" }
161-
listeners.value.forEach { it.onListChanged() }
162-
}
163-
164151
private fun notifyFeatureUpdated(oldFeature: T?, newFeature: T?) {
165-
logger.info { "Notifying listeners of feature update" }
166-
val featureKey = (oldFeature?.key ?: newFeature?.key) ?: return
152+
val featureKey = (oldFeature?.key ?: newFeature?.key) ?: run {
153+
logger.error { "Notification should have feature key, but none found" }
154+
return
155+
}
167156

157+
logger.info { "Notifying listeners on feature update" }
168158
listeners.value.forEach { it.onFeatureUpdated(featureKey) }
169159
}
170160
}

‎kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.UnsubscribeRequest
4343
import kotlinx.coroutines.CancellationException
4444
import kotlinx.coroutines.Deferred
4545
import kotlinx.serialization.json.JsonObject
46+
import kotlin.time.ExperimentalTime
4647

4748
private val logger = KotlinLogging.logger {}
4849

@@ -92,6 +93,7 @@ public open class Server(
9293

9394
private val sessionRegistry = ServerSessionRegistry()
9495

96+
@OptIn(ExperimentalTime::class)
9597
private val notificationService = FeatureNotificationService()
9698

9799
/**
@@ -111,17 +113,20 @@ public open class Server(
111113

112114
private val toolRegistry = FeatureRegistry<RegisteredTool>("Tool").apply {
113115
if (options.capabilities.tools?.listChanged ?: false) {
114-
addListener(notificationService.getToolFeatureListener())
116+
addListener(notificationService.getToolListChangedListener())
115117
}
116118
}
117119
private val promptRegistry = FeatureRegistry<RegisteredPrompt>("Prompt").apply {
118120
if (options.capabilities.prompts?.listChanged ?: false) {
119-
addListener(notificationService.getPromptFeatureListener())
121+
addListener(notificationService.getPromptListChangedListener())
120122
}
121123
}
122124
private val resourceRegistry = FeatureRegistry<RegisteredResource>("Resource").apply {
123125
if (options.capabilities.resources?.listChanged ?: false) {
124-
addListener(notificationService.geResourceFeatureListener())
126+
addListener(notificationService.getResourceListChangedListener())
127+
}
128+
if (options.capabilities.resources?.subscribe ?: false) {
129+
addListener(notificationService.getResourceUpdateListener())
125130
}
126131
}
127132

@@ -138,6 +143,7 @@ public open class Server(
138143

139144
public suspend fun close() {
140145
logger.debug { "Closing MCP server" }
146+
notificationService.close()
141147
sessions.forEach { (sessionId, session) ->
142148
logger.info { "Closing session $sessionId" }
143149
session.close()
@@ -215,14 +221,15 @@ public open class Server(
215221
// Register cleanup handler to remove session from list when it closes
216222
session.onClose {
217223
logger.debug { "Removing closed session from active sessions list" }
218-
// notificationService.unsubscribeFromListChangedNotification(session)
224+
notificationService.unsubscribeSession(session)
219225
sessionRegistry.removeSession(session.sessionId)
220226
}
227+
221228
logger.debug { "Server session connecting to transport" }
222229
session.connect(transport)
223230
logger.debug { "Server session successfully connected to transport" }
224231
sessionRegistry.addSession(session)
225-
// notificationService.subscribeToListChangedNotification(session)
232+
notificationService.subscribeSession(session)
226233

227234
_onConnect()
228235
return session
@@ -517,7 +524,7 @@ public open class Server(
517524
private suspend fun handleSubscribeResources(session: ServerSession, request: SubscribeRequest) {
518525
if (options.capabilities.resources?.subscribe ?: false) {
519526
logger.debug { "Subscribing to resources" }
520-
notificationService.subscribeToResourceUpdateNotifications(session, request.params.uri)
527+
notificationService.subscribeToResourceUpdate(session, request.params.uri)
521528
} else {
522529
logger.debug { "Failed to subscribe to resources: Server does not support resources capability" }
523530
}
@@ -526,7 +533,7 @@ public open class Server(
526533
private suspend fun handleUnsubscribeResources(session: ServerSession, request: UnsubscribeRequest) {
527534
if (options.capabilities.resources?.subscribe ?: false) {
528535
logger.debug { "Unsubscribing from resources" }
529-
notificationService.unsubscribeFromResourceUpdateNotifications(session, request.params.uri)
536+
notificationService.unsubscribeFromResourceUpdate(session, request.params.uri)
530537
} else {
531538
logger.debug { "Failed to unsubscribe from resources: Server does not support resources capability" }
532539
}

‎kotlin-sdk-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/InMemoryTransportTest.kt‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
11
package io.modelcontextprotocol.kotlin.sdk.integration
22

33
import io.modelcontextprotocol.kotlin.sdk.shared.InMemoryTransport
4+
import io.modelcontextprotocol.kotlin.sdk.types.BaseNotificationParams
45
import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification
56
import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage
7+
import io.modelcontextprotocol.kotlin.sdk.types.PromptListChangedNotification
8+
import io.modelcontextprotocol.kotlin.sdk.types.ResourceListChangedNotification
9+
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification
10+
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotificationParams
11+
import io.modelcontextprotocol.kotlin.sdk.types.ToolListChangedNotification
612
import io.modelcontextprotocol.kotlin.sdk.types.toJSON
713
import kotlinx.coroutines.test.runTest
814
import kotlin.test.BeforeTest
@@ -107,4 +113,94 @@ class InMemoryTransportTest {
107113
serverTransport.start()
108114
assertEquals(message, receivedMessage)
109115
}
116+
117+
@Test
118+
fun `should send ToolListChangedNotification from server to client`() = runTest {
119+
val notification = ToolListChangedNotification(
120+
BaseNotificationParams(),
121+
)
122+
123+
var receivedMessage: JSONRPCMessage? = null
124+
clientTransport.onMessage { msg ->
125+
receivedMessage = msg
126+
}
127+
128+
val rpcMessage = notification.toJSON()
129+
serverTransport.send(rpcMessage)
130+
assertEquals(rpcMessage, receivedMessage)
131+
}
132+
133+
@Test
134+
fun `should send PromptListChangedNotification from server to client`() = runTest {
135+
val notification = PromptListChangedNotification(
136+
BaseNotificationParams(),
137+
)
138+
139+
var receivedMessage: JSONRPCMessage? = null
140+
clientTransport.onMessage { msg ->
141+
receivedMessage = msg
142+
}
143+
144+
val rpcMessage = notification.toJSON()
145+
serverTransport.send(rpcMessage)
146+
assertEquals(rpcMessage, receivedMessage)
147+
}
148+
149+
@Test
150+
fun `should send ResourceListChangedNotification from server to client`() = runTest {
151+
val notification = ResourceListChangedNotification(
152+
BaseNotificationParams(),
153+
)
154+
155+
var receivedMessage: JSONRPCMessage? = null
156+
clientTransport.onMessage { msg ->
157+
receivedMessage = msg
158+
}
159+
160+
val rpcMessage = notification.toJSON()
161+
serverTransport.send(rpcMessage)
162+
assertEquals(rpcMessage, receivedMessage)
163+
}
164+
165+
@Test
166+
fun `should send ResourceUpdatedNotification from server to client`() = runTest {
167+
val notification = ResourceUpdatedNotification(
168+
ResourceUpdatedNotificationParams(
169+
uri = "file:///workspace/data.json",
170+
),
171+
)
172+
173+
var receivedMessage: JSONRPCMessage? = null
174+
clientTransport.onMessage { msg ->
175+
receivedMessage = msg
176+
}
177+
178+
val rpcMessage = notification.toJSON()
179+
serverTransport.send(rpcMessage)
180+
assertEquals(rpcMessage, receivedMessage)
181+
}
182+
183+
@Test
184+
fun `should handle multiple notifications in sequence`() = runTest {
185+
val notifications = listOf(
186+
ToolListChangedNotification(),
187+
PromptListChangedNotification(),
188+
ResourceListChangedNotification(),
189+
ResourceUpdatedNotification(ResourceUpdatedNotificationParams(uri = "file:///workspace/data.json")),
190+
)
191+
192+
val receivedMessages = mutableListOf<JSONRPCMessage>()
193+
clientTransport.onMessage { msg ->
194+
receivedMessages.add(msg)
195+
}
196+
197+
notifications.forEach { notification ->
198+
serverTransport.send(notification.toJSON())
199+
}
200+
201+
assertEquals(notifications.size, receivedMessages.size)
202+
notifications.forEachIndexed { index, notification ->
203+
assertEquals(notification.toJSON(), receivedMessages[index])
204+
}
205+
}
110206
}

0 commit comments

Comments
 (0)