diff --git a/kotlin-sdk-server/api/kotlin-sdk-server.api b/kotlin-sdk-server/api/kotlin-sdk-server.api index 8016d3bc..e5183f07 100644 --- a/kotlin-sdk-server/api/kotlin-sdk-server.api +++ b/kotlin-sdk-server/api/kotlin-sdk-server.api @@ -82,6 +82,8 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server { public final fun onConnect (Lkotlin/jvm/functions/Function0;)V public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V public final fun ping (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public final fun removeNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V + public final fun removeNotificationHandler (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/Method;)V public final fun removePrompt (Ljava/lang/String;)Z public final fun removePrompts (Ljava/util/List;)I public final fun removeResource (Ljava/lang/String;)Z @@ -93,6 +95,8 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server { public final fun sendResourceListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun sendResourceUpdated (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/ResourceUpdatedNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public final fun sendToolListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public final fun setNotificationHandler (Lio/modelcontextprotocol/kotlin/sdk/types/Method;Lkotlin/jvm/functions/Function1;)V + public final fun setNotificationHandler (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/Method;Lkotlin/jvm/functions/Function1;)V } public final class io/modelcontextprotocol/kotlin/sdk/server/ServerOptions : io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions { diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Features.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Feature.kt similarity index 100% rename from kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Features.kt rename to kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Feature.kt diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureNotificationService.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureNotificationService.kt new file mode 100644 index 00000000..397e6e19 --- /dev/null +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureNotificationService.kt @@ -0,0 +1,373 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.github.oshai.kotlinlogging.KotlinLogging +import io.modelcontextprotocol.kotlin.sdk.types.Notification +import io.modelcontextprotocol.kotlin.sdk.types.PromptListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotificationParams +import io.modelcontextprotocol.kotlin.sdk.types.ServerNotification +import io.modelcontextprotocol.kotlin.sdk.types.ToolListChangedNotification +import kotlinx.atomicfu.atomic +import kotlinx.atomicfu.getAndUpdate +import kotlinx.collections.immutable.persistentMapOf +import kotlinx.collections.immutable.persistentSetOf +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.cancel +import kotlinx.coroutines.channels.BufferOverflow +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.flow.MutableSharedFlow +import kotlinx.coroutines.flow.SharedFlow +import kotlinx.coroutines.flow.takeWhile +import kotlinx.coroutines.joinAll +import kotlinx.coroutines.launch +import kotlin.time.Clock +import kotlin.time.ExperimentalTime + +/** + * Represents an event for the notification service. + * + * @property timestamp A timestamp for the event. + */ +private sealed class NotificationEvent(open val timestamp: Long) + +/** + * Represents an event for a notification. + * + * @property notification The notification associated with the event. + */ +private class SendEvent(override val timestamp: Long, val notification: Notification) : NotificationEvent(timestamp) + +/** Represents an event marking the end of notification processing. */ +private class EndEvent(override val timestamp: Long) : NotificationEvent(timestamp) + +/** + * Represents a job that handles session-specific notifications, processing events + * and delivering relevant notifications to the associated session. + * + * This class listens to a stream of notification events and processes them + * based on the event type and the resource subscriptions associated with the session. + * It allows subscribing to or unsubscribing from specific resource keys for granular + * notification handling. The job can also be canceled to stop processing further events. + * Notification with timestamps older than the starting timestamp are skipped. + */ +private class SessionNotificationJob { + private val job: Job + private val resourceSubscriptions = atomic(persistentMapOf()) + private val logger = KotlinLogging.logger {} + + /** + * Constructor for the SessionNotificationJob, responsible for processing notification events + * and dispatching appropriate notifications to the provided server session. The job operates + * within the given coroutine scope and begins handling events starting from the specified + * timestamp. + * + * @param session The server session where notifications will be dispatched. + * @param scope The coroutine scope in which this job operates. + * @param events A shared flow of notification events that the job listens to. + * @param fromTimestamp The timestamp from which the job starts processing events. + */ + constructor( + session: ServerSession, + scope: CoroutineScope, + events: SharedFlow, + fromTimestamp: Long, + ) { + logger.info { "Starting notification job from timestamp $fromTimestamp for sessionId: ${session.sessionId} " } + job = scope.launch { + events.takeWhile { it !is EndEvent }.collect { event -> + when (event) { + is SendEvent -> { + handleSendNotificationEvent(event, session, fromTimestamp) + } + + else -> { + logger.warn { "Skipping event: $event" } + } + } + } + } + } + + /** + * Handles sending a notification event to a specific server session. + * + * @param event The notification event to be processed. + * @param session The server session to which the notification should be sent. + * @param fromTimestamp The timestamp to filter events. + * Notifications with timestamps older than this value are skipped. + */ + private suspend fun handleSendNotificationEvent(event: SendEvent, session: ServerSession, fromTimestamp: Long) { + if (event.timestamp < fromTimestamp) { + logger.info { + "Skipping event with id: ${event.timestamp} as it is older than startingEventId $fromTimestamp: $event" + } + return + } + when (val notification = event.notification) { + is PromptListChangedNotification, + is ResourceListChangedNotification, + is ToolListChangedNotification, + -> { + logger.info { "Sending list changed notification for sessionId: ${session.sessionId}" } + session.notification(notification) + } + + is ResourceUpdatedNotification -> { + resourceSubscriptions.value[notification.params.uri]?.let { resourceFromTimestamp -> + if (event.timestamp >= resourceFromTimestamp) { + logger.info { + "Sending resource updated notification for resource ${notification.params.uri} " + + "to sessionId: ${session.sessionId}" + } + session.notification(notification) + } else { + logger.info { + "Skipping resource updated notification for resource ${notification.params.uri} " + + "as it is older than subscription timestamp $resourceFromTimestamp" + } + } + } ?: { + logger.info { + "No subscription for resource ${notification.params.uri}. " + + "Skipping notification: $notification" + } + } + } + + else -> { + logger.warn { "Skipping notification: $notification" } + } + } + } + + /** + * Subscribes to a resource identified by the given feature key. + * + * @param resourceKey The key representing the resource to subscribe to. + * @param timestamp The timestamp of the subscription. + */ + fun subscribe(resourceKey: FeatureKey, timestamp: Long) { + resourceSubscriptions.getAndUpdate { it.put(resourceKey, timestamp) } + } + + /** + * Unsubscribes from a resource identified by the given feature key. + * + * @param resourceKey The key representing the resource to unsubscribe from. + */ + fun unsubscribe(resourceKey: FeatureKey) { + resourceSubscriptions.getAndUpdate { it.remove(resourceKey) } + } + + /** + * Waits for the notification service to complete its operations. + */ + suspend fun join() { + job.join() + } + + /** + * Cancels the notification service job. + */ + fun cancel() { + job.cancel() + } +} + +/** + * Service responsible for managing and emitting notifications related to feature changes. + * + * This service facilitates notification subscriptions for different sessions and supports managing + * listeners for feature-related events. Notifications include changes in tool lists, prompt lists, + * resource lists, and updates to specific resources. + * + * This class operates on a background coroutine scope to handle notifications asynchronously. + * It maintains jobs associated with sessions and features for controlling active subscriptions. + * + * Key Responsibilities: + * - Emit notifications for various feature-related events. + * - Provide listeners for handling feature change events. + * - Allow clients to subscribe or unsubscribe from specific notifications. + * + * Notifications managed: + * - Tool list change notifications. + * - Prompt list change notifications. + * - Resource list change notifications. + * - Resource updates pertaining to specific resources. + */ +internal class FeatureNotificationService( + notificationBufferCapacity: Int = Channel.UNLIMITED, + @OptIn(ExperimentalTime::class) + private val clock: Clock = Clock.System, +) { + /** Coroutine scope used to handle asynchronous notifications. */ + private val notificationScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + + /** Shared flow used to emit events within the feature notification service. */ + private val notificationEvents = MutableSharedFlow( + extraBufferCapacity = notificationBufferCapacity, + onBufferOverflow = BufferOverflow.SUSPEND, + ) + + /** Active emit jobs. */ + private val emitJobs = atomic(persistentSetOf()) + + /** Notification jobs associated with sessions. */ + private val sessionNotificationJobs = atomic(persistentMapOf()) + + private val logger = KotlinLogging.logger {} + + private fun featureListener(notificationProvider: (FeatureKey) -> ServerNotification): FeatureListener = + object : FeatureListener { + override fun onFeatureUpdated(featureKey: FeatureKey) { + val notification = notificationProvider(featureKey) + logger.debug { "Emitting notification: ${notification.method.value}" } + emit(notification) + } + } + + /** Listener for tool feature events. */ + internal val toolListChangedListener: FeatureListener = featureListener { ToolListChangedNotification() } + + /** Listener for prompt feature events. */ + internal val promptListChangedListener: FeatureListener = featureListener { PromptListChangedNotification() } + + /** Listener for resource feature events. */ + internal val resourceListChangedListener: FeatureListener = featureListener { ResourceListChangedNotification() } + + /** Listener for resource update events. */ + internal val resourceUpdatedListener: FeatureListener = + featureListener { ResourceUpdatedNotification(ResourceUpdatedNotificationParams(uri = it)) } + + /** + * Subscribes session to list changed notifications for all features and resource update notifications. + * For each session the job is created and stored until the [unsubscribeSession] method is called. + * In case of session already subscribed to list changed notifications, the method will skip the subscription and + * continue to send notification using the existing job. + * + * @param session The session to subscribe. + */ + internal fun subscribeSession(session: ServerSession) { + logger.info { "Subscribing session for notifications sessionId: ${session.sessionId}" } + + val timestamp = getCurrentTimestamp() + + sessionNotificationJobs.getAndUpdate { + if (it.containsKey(session.sessionId)) { + logger.info { "Session already subscribed: ${session.sessionId}" } + return@getAndUpdate it + } else { + it.put( + session.sessionId, + SessionNotificationJob( + session = session, + scope = notificationScope, + events = notificationEvents, + // Save the first event id to process, as notification can be emitted after the subscription + fromTimestamp = timestamp, + ), + ) + } + } + + logger.info { "Subscribed session for notifications sessionId: ${session.sessionId}" } + } + + /** + * Unsubscribes a session from list changed notifications for all features. + * Cancels and removes the job associated with the given session's notifications. + * + * @param session The session to unsubscribe from list changed notifications. + */ + internal fun unsubscribeSession(session: ServerSession) { + logger.info { "Unsubscribing from list changed notifications for sessionId: ${session.sessionId}" } + sessionNotificationJobs.getAndUpdate { + it[session.sessionId]?.cancel() + it.remove(session.sessionId) + } + logger.info { "Unsubscribed from list changed notifications for sessionId: ${session.sessionId}" } + } + + /** + * Subscribes a session to notifications for resource updates pertaining to the given resource key. + * + * @param session The session to subscribe. + * @param resourceKey The resource key to subscribe to. + */ + internal fun subscribeToResourceUpdate(session: ServerSession, resourceKey: FeatureKey) { + logger.info { "Subscribing to resource $resourceKey update notifications for sessionId: ${session.sessionId}" } + // Set starting event id for resources notifications to skip events emitted before the subscription + sessionNotificationJobs.value[session.sessionId]?.subscribe(resourceKey, getCurrentTimestamp()) + logger.info { "Subscribed to resource $resourceKey update notifications for sessionId: ${session.sessionId}" } + } + + /** + * Unsubscribes a session from notifications for resource updates pertaining to the given resource key. + * + * @param session The session to unsubscribe from. + * @param resourceKey The resource key to unsubscribe from. + */ + internal fun unsubscribeFromResourceUpdate(session: ServerSession, resourceKey: FeatureKey) { + logger.info { + "Unsubscribing from resource $resourceKey update notifications for sessionId: ${session.sessionId}" + } + sessionNotificationJobs.value[session.sessionId]?.unsubscribe(resourceKey) + logger.info { + "Unsubscribed from resource $resourceKey update notifications for sessionId: ${session.sessionId}" + } + } + + /** Emits a notification to all active sessions. */ + private fun emit(notification: Notification) { + // Create a timestamp before emit to ensure notifications are processed in order + val timestamp = getCurrentTimestamp() + + logger.info { "Emitting notification $timestamp: $notification" } + + // Launching emit lazily to put it to the jobs queue before the completion + val job = notificationScope.launch(start = CoroutineStart.LAZY) { + logger.debug { "Actually emitting notification $timestamp: $notification" } + notificationEvents.emit(SendEvent(timestamp, notification)) + logger.debug { "Notification emitted $timestamp: $notification" } + } + + // Add job to set before starting + emitJobs.getAndUpdate { it.add(job) } + + // Register completion + job.invokeOnCompletion { + emitJobs.getAndUpdate { it.remove(job) } + } + + // Start the job after it's safely added + job.start() + } + + /** Returns the current timestamp in milliseconds. */ + @OptIn(ExperimentalTime::class) + private fun getCurrentTimestamp(): Long = clock.now().toEpochMilliseconds() + + suspend fun close() { + logger.info { "Closing feature notification service" } + + // Making sure all emitting jobs are completed + emitJobs.value.joinAll() + + // Emitting end event to complete all session notification jobs + notificationScope.launch { + logger.info { "Emitting end event" } + notificationEvents.emit(EndEvent(getCurrentTimestamp())) + logger.info { "End event emitted" } + }.join() + + // Making sure all session notification jobs are completed (after receiving the end event) + sessionNotificationJobs.value.values.forEach { it.join() } + // Cancelling notification scope to stop processing further events + notificationScope.cancel() + } +} diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureRegistry.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureRegistry.kt index 1749ef64..0ed65588 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureRegistry.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/FeatureRegistry.kt @@ -5,9 +5,17 @@ import kotlinx.atomicfu.atomic import kotlinx.atomicfu.getAndUpdate import kotlinx.atomicfu.update import kotlinx.collections.immutable.minus +import kotlinx.collections.immutable.persistentListOf import kotlinx.collections.immutable.persistentMapOf import kotlinx.collections.immutable.toPersistentSet +/** + * A listener interface for receiving notifications about feature changes in registry. + */ +internal interface FeatureListener { + fun onFeatureUpdated(featureKey: String) +} + /** * A generic registry for managing features of a specified type. This class provides thread-safe * operations for adding, removing, and retrieving features from the registry. @@ -33,6 +41,16 @@ internal class FeatureRegistry(private val featureType: String) { internal val values: Map get() = registry.value + private val listeners = atomic(persistentListOf()) + + internal fun addListener(listener: FeatureListener) { + listeners.update { it.add(listener) } + } + + internal fun removeListener(listener: FeatureListener) { + listeners.update { it.remove(listener) } + } + /** * Adds the specified feature to the registry. * @@ -40,8 +58,11 @@ internal class FeatureRegistry(private val featureType: String) { */ internal fun add(feature: T) { logger.info { "Adding $featureType: \"${feature.key}\"" } - registry.update { current -> current.put(feature.key, feature) } + val oldMap = registry.getAndUpdate { current -> current.put(feature.key, feature) } + val oldFeature = oldMap[feature.key] + logger.info { "Added $featureType: \"${feature.key}\"" } + notifyFeatureUpdated(oldFeature, feature) } /** @@ -52,8 +73,13 @@ internal class FeatureRegistry(private val featureType: String) { */ internal fun addAll(features: List) { logger.info { "Adding ${featureType}s: ${features.size}" } - registry.update { current -> current.putAll(features.associateBy { it.key }) } + val oldMap = registry.getAndUpdate { current -> current.putAll(features.associateBy { it.key }) } + logger.info { "Added ${featureType}s: ${features.size}" } + for (feature in features) { + val oldFeature = oldMap[feature.key] + notifyFeatureUpdated(oldFeature, feature) + } } /** @@ -66,15 +92,17 @@ internal class FeatureRegistry(private val featureType: String) { logger.info { "Removing $featureType: \"$key\"" } val oldMap = registry.getAndUpdate { current -> current.remove(key) } - val removed = key in oldMap - logger.info { - if (removed) { - "Removed $featureType: \"$key\"" - } else { - "$featureType not found: \"$key\"" - } + val removedFeature = oldMap[key] + val removed = removedFeature != null + + if (removed) { + logger.info { "Removed $featureType: \"$key\"" } + notifyFeatureUpdated(removedFeature, null) + } else { + logger.info { "$featureType not found: \"$key\"" } } - return key in oldMap + + return removed } /** @@ -87,13 +115,16 @@ internal class FeatureRegistry(private val featureType: String) { logger.info { "Removing ${featureType}s: ${keys.size}" } val oldMap = registry.getAndUpdate { current -> current - keys.toPersistentSet() } - val removedCount = keys.count { it in oldMap } - logger.info { - if (removedCount > 0) { - "Removed ${featureType}s: $removedCount" - } else { - "No $featureType were removed" - } + val removedFeatures = keys.mapNotNull { oldMap[it] } + val removedCount = removedFeatures.size + + if (removedCount > 0) { + logger.info { "Removed ${featureType}s: $removedCount" } + } else { + logger.info { "No $featureType were removed" } + } + removedFeatures.forEach { + notifyFeatureUpdated(it, null) } return removedCount @@ -108,13 +139,22 @@ internal class FeatureRegistry(private val featureType: String) { internal fun get(key: FeatureKey): T? { logger.info { "Getting $featureType: \"$key\"" } val feature = registry.value[key] - logger.info { - if (feature != null) { - "Got $featureType: \"$key\"" - } else { - "$featureType not found: \"$key\"" - } + if (feature != null) { + logger.info { "Got $featureType: \"$key\"" } + } else { + logger.info { "$featureType not found: \"$key\"" } } + return feature } + + private fun notifyFeatureUpdated(oldFeature: T?, newFeature: T?) { + val featureKey = (oldFeature?.key ?: newFeature?.key) ?: run { + logger.error { "Notification should have feature key, but none found" } + return + } + + logger.info { "Notifying listeners on feature update" } + listeners.value.forEach { it.onFeatureUpdated(featureKey) } + } } diff --git a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt index 7fa7a98f..ef52d728 100644 --- a/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt +++ b/kotlin-sdk-server/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/Server.kt @@ -26,6 +26,7 @@ import io.modelcontextprotocol.kotlin.sdk.types.ListToolsRequest import io.modelcontextprotocol.kotlin.sdk.types.ListToolsResult import io.modelcontextprotocol.kotlin.sdk.types.LoggingMessageNotification import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.Notification import io.modelcontextprotocol.kotlin.sdk.types.Prompt import io.modelcontextprotocol.kotlin.sdk.types.PromptArgument import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequest @@ -33,12 +34,16 @@ import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceResult import io.modelcontextprotocol.kotlin.sdk.types.Resource import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.SubscribeRequest import io.modelcontextprotocol.kotlin.sdk.types.TextContent import io.modelcontextprotocol.kotlin.sdk.types.Tool import io.modelcontextprotocol.kotlin.sdk.types.ToolAnnotations import io.modelcontextprotocol.kotlin.sdk.types.ToolSchema +import io.modelcontextprotocol.kotlin.sdk.types.UnsubscribeRequest import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.Deferred import kotlinx.serialization.json.JsonObject +import kotlin.time.ExperimentalTime private val logger = KotlinLogging.logger {} @@ -58,6 +63,11 @@ public class ServerOptions(public val capabilities: ServerCapabilities, enforceS * You can register tools, prompts, and resources using [addTool], [addPrompt], and [addResource]. * The server will then automatically handle listing and retrieval requests from the client. * + * In case the server supports feature list notification or resource substitution, + * the server will automatically send notifications for all connected clients. + * Currently, after subscription to a resource, the server will NOT send the subscription confirmation + * as this response schema is not defined in the protocol. + * * @param serverInfo Information about this server implementation (name, version). * @param options Configuration options for the server. * @param instructionsProvider Optional provider for instructions from the server to the client about how to use @@ -86,14 +96,6 @@ public open class Server( block: Server.() -> Unit = {}, ) : this(serverInfo, options, { instructions }, block) - private val sessionRegistry = ServerSessionRegistry() - - /** - * Provides a snapshot of all sessions currently registered in the server - */ - public val sessions: Map - get() = sessionRegistry.sessions - @Suppress("ktlint:standard:backing-property-naming") private var _onInitialized: (() -> Unit) = {} @@ -103,14 +105,51 @@ public open class Server( @Suppress("ktlint:standard:backing-property-naming") private var _onClose: () -> Unit = {} - private val toolRegistry = FeatureRegistry("Tool") - private val promptRegistry = FeatureRegistry("Prompt") - private val resourceRegistry = FeatureRegistry("Resource") + @OptIn(ExperimentalTime::class) + private val notificationService = FeatureNotificationService() + + private val sessionRegistry = ServerSessionRegistry() + private val toolRegistry = FeatureRegistry("Tool").apply { + if (options.capabilities.tools?.listChanged ?: false) { + addListener(notificationService.toolListChangedListener) + } + } + private val promptRegistry = FeatureRegistry("Prompt").apply { + if (options.capabilities.prompts?.listChanged ?: false) { + addListener(notificationService.promptListChangedListener) + } + } + private val resourceRegistry = FeatureRegistry("Resource").apply { + if (options.capabilities.resources?.listChanged ?: false) { + addListener(notificationService.resourceListChangedListener) + } + if (options.capabilities.resources?.subscribe ?: false) { + addListener(notificationService.resourceUpdatedListener) + } + } + + /** + * Provides a snapshot of all sessions currently registered in the server + */ + public val sessions: Map + get() = sessionRegistry.sessions + + /** + * Provides a snapshot of all tools currently registered in the server + */ public val tools: Map get() = toolRegistry.values + + /** + * Provides a snapshot of all prompts currently registered in the server + */ public val prompts: Map get() = promptRegistry.values + + /** + * Provides a snapshot of all resources currently registered in the server + */ public val resources: Map get() = resourceRegistry.values @@ -120,6 +159,7 @@ public open class Server( public suspend fun close() { logger.debug { "Closing MCP server" } + notificationService.close() sessions.forEach { (sessionId, session) -> logger.info { "Closing session $sessionId" } session.close() @@ -182,17 +222,32 @@ public open class Server( session.setRequestHandler(Method.Defined.ResourcesTemplatesList) { _, _ -> handleListResourceTemplates() } + if (options.capabilities.resources?.subscribe ?: false) { + session.setRequestHandler(Method.Defined.ResourcesSubscribe) { request, _ -> + handleSubscribeResources(session, request) + // Does not return any confirmation as the structure is not stated in the protocol + null + } + session.setRequestHandler(Method.Defined.ResourcesUnsubscribe) { request, _ -> + handleUnsubscribeResources(session, request) + // Does not return any confirmation as the structure is not stated in the protocol + null + } + } } // Register cleanup handler to remove session from list when it closes session.onClose { logger.debug { "Removing closed session from active sessions list" } + notificationService.unsubscribeSession(session) sessionRegistry.removeSession(session.sessionId) } + logger.debug { "Server session connecting to transport" } session.connect(transport) logger.debug { "Server session successfully connected to transport" } sessionRegistry.addSession(session) + notificationService.subscribeSession(session) _onConnect() return session @@ -484,7 +539,25 @@ public open class Server( } // --- Internal Handlers --- - private suspend fun handleListTools(): ListToolsResult { + private fun handleSubscribeResources(session: ServerSession, request: SubscribeRequest) { + if (options.capabilities.resources?.subscribe ?: false) { + logger.debug { "Subscribing to resources" } + notificationService.subscribeToResourceUpdate(session, request.params.uri) + } else { + logger.debug { "Failed to subscribe to resources: Server does not support resources capability" } + } + } + + private fun handleUnsubscribeResources(session: ServerSession, request: UnsubscribeRequest) { + if (options.capabilities.resources?.subscribe ?: false) { + logger.debug { "Unsubscribing from resources" } + notificationService.unsubscribeFromResourceUpdate(session, request.params.uri) + } else { + logger.debug { "Failed to unsubscribe from resources: Server does not support resources capability" } + } + } + + private fun handleListTools(): ListToolsResult { val toolList = tools.values.map { it.tool } return ListToolsResult(tools = toolList, nextCursor = null) } @@ -518,7 +591,7 @@ public open class Server( } } - private suspend fun handleListPrompts(): ListPromptsResult { + private fun handleListPrompts(): ListPromptsResult { logger.debug { "Handling list prompts request" } return ListPromptsResult(prompts = prompts.values.map { it.prompt }) } @@ -534,7 +607,7 @@ public open class Server( return prompt.messageProvider(request) } - private suspend fun handleListResources(): ListResourcesResult { + private fun handleListResources(): ListResourcesResult { logger.debug { "Handling list resources request" } return ListResourcesResult(resources = resources.values.map { it.resource }) } @@ -550,7 +623,7 @@ public open class Server( return resource.readHandler(request) } - private suspend fun handleListResourceTemplates(): ListResourceTemplatesResult { + private fun handleListResourceTemplates(): ListResourceTemplatesResult { // If you have resource templates, return them here. For now, return empty. return ListResourceTemplatesResult(listOf()) } @@ -675,4 +748,30 @@ public open class Server( } } // End the ServerSession redirection section + + // Start the notification handling section + public fun setNotificationHandler(method: Method, handler: (notification: T) -> Deferred) { + sessions.forEach { (_, session) -> + session.setNotificationHandler(method, handler) + } + } + + public fun removeNotificationHandler(method: Method) { + sessions.forEach { (_, session) -> + session.removeNotificationHandler(method) + } + } + + public fun setNotificationHandler( + sessionId: String, + method: Method, + handler: (notification: T) -> Deferred, + ) { + sessionRegistry.getSessionOrNull(sessionId)?.setNotificationHandler(method, handler) + } + + public fun removeNotificationHandler(sessionId: String, method: Method) { + sessionRegistry.getSessionOrNull(sessionId)?.removeNotificationHandler(method) + } + // End the notification handling section } diff --git a/kotlin-sdk-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/InMemoryTransportTest.kt b/kotlin-sdk-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/InMemoryTransportTest.kt index d13f5848..83e56184 100644 --- a/kotlin-sdk-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/InMemoryTransportTest.kt +++ b/kotlin-sdk-test/src/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/integration/InMemoryTransportTest.kt @@ -1,8 +1,14 @@ package io.modelcontextprotocol.kotlin.sdk.integration import io.modelcontextprotocol.kotlin.sdk.shared.InMemoryTransport +import io.modelcontextprotocol.kotlin.sdk.types.BaseNotificationParams import io.modelcontextprotocol.kotlin.sdk.types.InitializedNotification import io.modelcontextprotocol.kotlin.sdk.types.JSONRPCMessage +import io.modelcontextprotocol.kotlin.sdk.types.PromptListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotificationParams +import io.modelcontextprotocol.kotlin.sdk.types.ToolListChangedNotification import io.modelcontextprotocol.kotlin.sdk.types.toJSON import kotlinx.coroutines.test.runTest import kotlin.test.BeforeTest @@ -107,4 +113,94 @@ class InMemoryTransportTest { serverTransport.start() assertEquals(message, receivedMessage) } + + @Test + fun `should send ToolListChangedNotification from server to client`() = runTest { + val notification = ToolListChangedNotification( + BaseNotificationParams(), + ) + + var receivedMessage: JSONRPCMessage? = null + clientTransport.onMessage { msg -> + receivedMessage = msg + } + + val rpcMessage = notification.toJSON() + serverTransport.send(rpcMessage) + assertEquals(rpcMessage, receivedMessage) + } + + @Test + fun `should send PromptListChangedNotification from server to client`() = runTest { + val notification = PromptListChangedNotification( + BaseNotificationParams(), + ) + + var receivedMessage: JSONRPCMessage? = null + clientTransport.onMessage { msg -> + receivedMessage = msg + } + + val rpcMessage = notification.toJSON() + serverTransport.send(rpcMessage) + assertEquals(rpcMessage, receivedMessage) + } + + @Test + fun `should send ResourceListChangedNotification from server to client`() = runTest { + val notification = ResourceListChangedNotification( + BaseNotificationParams(), + ) + + var receivedMessage: JSONRPCMessage? = null + clientTransport.onMessage { msg -> + receivedMessage = msg + } + + val rpcMessage = notification.toJSON() + serverTransport.send(rpcMessage) + assertEquals(rpcMessage, receivedMessage) + } + + @Test + fun `should send ResourceUpdatedNotification from server to client`() = runTest { + val notification = ResourceUpdatedNotification( + ResourceUpdatedNotificationParams( + uri = "file:///workspace/data.json", + ), + ) + + var receivedMessage: JSONRPCMessage? = null + clientTransport.onMessage { msg -> + receivedMessage = msg + } + + val rpcMessage = notification.toJSON() + serverTransport.send(rpcMessage) + assertEquals(rpcMessage, receivedMessage) + } + + @Test + fun `should handle multiple notifications in sequence`() = runTest { + val notifications = listOf( + ToolListChangedNotification(), + PromptListChangedNotification(), + ResourceListChangedNotification(), + ResourceUpdatedNotification(ResourceUpdatedNotificationParams(uri = "file:///workspace/data.json")), + ) + + val receivedMessages = mutableListOf() + clientTransport.onMessage { msg -> + receivedMessages.add(msg) + } + + notifications.forEach { notification -> + serverTransport.send(notification.toJSON()) + } + + assertEquals(notifications.size, receivedMessages.size) + notifications.forEachIndexed { index, notification -> + assertEquals(notification.toJSON(), receivedMessages[index]) + } + } } diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsNotificationTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsNotificationTest.kt new file mode 100644 index 00000000..e2cfe294 --- /dev/null +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsNotificationTest.kt @@ -0,0 +1,116 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.modelcontextprotocol.kotlin.sdk.types.GetPromptResult +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.Prompt +import io.modelcontextprotocol.kotlin.sdk.types.PromptListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.test.runTest +import org.awaitility.kotlin.await +import org.awaitility.kotlin.untilAsserted +import org.junit.jupiter.api.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +class ServerPromptsNotificationTest : AbstractServerFeaturesTest() { + + override fun getServerCapabilities(): ServerCapabilities = ServerCapabilities( + prompts = ServerCapabilities.Prompts(true), + ) + + @Test + fun `addPrompt should send notification`() = runTest { + // Configure notification handler + var promptListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + promptListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Add a prompt + val testPrompt = Prompt("test-prompt", "Test Prompt", null) + server.addPrompt(testPrompt) { + GetPromptResult( + description = "Test prompt description", + messages = listOf(), + ) + } + + // Remove the prompt + val result = server.removePrompt(testPrompt.name) + + // Verify the prompt was removed + assertTrue(result, "Prompt should be removed successfully") + + // Verify that the notification was sent + await untilAsserted { + assertTrue(promptListChangedNotificationReceived, "Notification should be sent when prompt is added") + } + } + + @Test + fun `removePrompts should remove multiple prompts and send two notifications`() = runTest { + // Configure notification handler + var promptListChangedNotificationReceivedCount = 0 + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + promptListChangedNotificationReceivedCount += 1 + CompletableDeferred(Unit) + } + + // Add prompts + val testPrompt1 = Prompt("test-prompt-1", "Test Prompt 1", null) + val testPrompt2 = Prompt("test-prompt-2", "Test Prompt 2", null) + + server.addPrompt(testPrompt1) { + GetPromptResult( + description = "Test prompt description 1", + messages = listOf(), + ) + } + server.addPrompt(testPrompt2) { + GetPromptResult( + description = "Test prompt description 2", + messages = listOf(), + ) + } + + // Remove the prompts + val result = server.removePrompts(listOf(testPrompt1.name, testPrompt2.name)) + + // Verify the prompts were removed + assertEquals(2, result, "Both prompts should be removed") + + // Verify that the notifications were sent twice + await untilAsserted { + assertEquals( + 4, + promptListChangedNotificationReceivedCount, + "Two notifications should be sent when prompts are added and two when removed", + ) + } + } + + @Test + fun `notification should not be send when removed prompt does not exists`() = runTest { + // Track notifications + var promptListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + promptListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Try to remove a non-existent prompt + val result = server.removePrompt("non-existent-prompt") + + // Verify the result + assertFalse(result, "Removing non-existent prompt should return false") + await untilAsserted { + assertFalse( + promptListChangedNotificationReceived, + "No notification should be sent when prompt doesn't exist", + ) + } + } +} diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsTest.kt index 60de0cab..5f49d04c 100644 --- a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsTest.kt +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerPromptsTest.kt @@ -17,11 +17,18 @@ import kotlin.test.assertTrue class ServerPromptsTest : AbstractServerFeaturesTest() { override fun getServerCapabilities(): ServerCapabilities = ServerCapabilities( - prompts = ServerCapabilities.Prompts(false), + prompts = ServerCapabilities.Prompts(null), ) @Test - fun `removePrompt should remove a prompt`() = runTest { + fun `removePrompt should remove a prompt and do not send notification`() = runTest { + // Configure notification handler + var promptListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + promptListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + // Add a prompt val testPrompt = Prompt("test-prompt", "Test Prompt", null) server.addPrompt(testPrompt) { @@ -36,13 +43,30 @@ class ServerPromptsTest : AbstractServerFeaturesTest() { // Verify the prompt was removed assertTrue(result, "Prompt should be removed successfully") + + // Verify that the notification was not sent + assertFalse( + promptListChangedNotificationReceived, + "No notification should be sent when prompts capability is not supported", + ) } @Test - fun `removePrompts should remove multiple prompts and send notification`() = runTest { + fun `removePrompts should remove multiple prompts`() = runTest { + // Configure notification handler + var promptListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + promptListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + // Add prompts val testPrompt1 = Prompt("test-prompt-1", "Test Prompt 1", null) val testPrompt2 = Prompt("test-prompt-2", "Test Prompt 2", null) + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + throw IllegalStateException("Notification should not be sent") + } + server.addPrompt(testPrompt1) { GetPromptResult( description = "Test prompt description 1", @@ -61,6 +85,12 @@ class ServerPromptsTest : AbstractServerFeaturesTest() { // Verify the prompts were removed assertEquals(2, result, "Both prompts should be removed") + + // Verify that the notification was not sent + assertFalse( + promptListChangedNotificationReceived, + "No notification should be sent when prompts capability is not supported", + ) } @Test @@ -82,6 +112,12 @@ class ServerPromptsTest : AbstractServerFeaturesTest() { @Test fun `removePrompt should throw when prompts capability is not supported`() = runTest { + var promptListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsPromptsListChanged) { + promptListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + // Create server without prompts capability val serverOptions = ServerOptions( capabilities = ServerCapabilities(), @@ -96,5 +132,11 @@ class ServerPromptsTest : AbstractServerFeaturesTest() { server.removePrompt("test-prompt") } assertEquals("Server does not support prompts capability.", exception.message) + + // Verify that the notification was not sent + assertFalse( + promptListChangedNotificationReceived, + "No notification should be sent when prompts capability is not supported", + ) } } diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationSubscribeTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationSubscribeTest.kt new file mode 100644 index 00000000..62844501 --- /dev/null +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationSubscribeTest.kt @@ -0,0 +1,186 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceResult +import io.modelcontextprotocol.kotlin.sdk.types.ResourceListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.SubscribeRequest +import io.modelcontextprotocol.kotlin.sdk.types.SubscribeRequestParams +import io.modelcontextprotocol.kotlin.sdk.types.TextResourceContents +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.test.runTest +import org.awaitility.kotlin.await +import org.awaitility.kotlin.untilAsserted +import org.junit.jupiter.api.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +class ServerResourcesNotificationSubscribeTest : AbstractServerFeaturesTest() { + + override fun getServerCapabilities(): ServerCapabilities = ServerCapabilities( + resources = ServerCapabilities.Resources(null, true), + ) + + @Test + fun `removeResource should send resource update notification`() = runTest { + val notifications = mutableListOf() + client.setNotificationHandler(Method.Defined.NotificationsResourcesUpdated) { + notifications.add(it) + CompletableDeferred(Unit) + } + + // Add resources + val testResourceUri1 = "test://resource1" + val testResourceUri2 = "test://resource2" + + server.addResource( + uri = testResourceUri1, + name = "Test Resource 1", + description = "A test resource 1", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content 1", + uri = testResourceUri1, + mimeType = "text/plain", + ), + ), + ) + } + + client.subscribeResource(SubscribeRequest(SubscribeRequestParams(uri = testResourceUri1))) + + server.addResource( + uri = testResourceUri2, + name = "Test Resource 2", + description = "A test resource 2", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content 2", + uri = testResourceUri2, + mimeType = "text/plain", + ), + ), + ) + } + + // Remove the resource + val result = server.removeResource(testResourceUri1) + + // Verify the resource was removed + assertTrue(result, "Resource should be removed successfully") + + // Verify that the notification was sent + await untilAsserted { + assertEquals(1, notifications.size, "Notification should be sent when resource 1 was deleted") + } + assertEquals(testResourceUri1, notifications[0].params.uri, "Notification should contain the resource 1 URI") + } + + @Test + fun `removeResource for two resources should send two separate notifications`() = runTest { + val notifications = mutableListOf() + client.setNotificationHandler(Method.Defined.NotificationsResourcesUpdated) { + notifications.add(it) + CompletableDeferred(Unit) + } + + // Add resources + val testResourceUri1 = "test://resource1" + val testResourceUri2 = "test://resource2" + + server.addResource( + uri = testResourceUri1, + name = "Test Resource 1", + description = "A test resource 1", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content 1", + uri = testResourceUri1, + mimeType = "text/plain", + ), + ), + ) + } + + server.addResource( + uri = testResourceUri2, + name = "Test Resource 2", + description = "A test resource 2", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content 2", + uri = testResourceUri2, + mimeType = "text/plain", + ), + ), + ) + } + + client.subscribeResource(SubscribeRequest(SubscribeRequestParams(uri = testResourceUri1))) + client.subscribeResource(SubscribeRequest(SubscribeRequestParams(uri = testResourceUri2))) + + // Remove the resource + val result1 = server.removeResource(testResourceUri1) + val result2 = server.removeResource(testResourceUri2) + + // Verify the resource was removed + assertTrue(result1, "Resource 1 should be removed successfully") + assertTrue(result2, "Resource 2 should be removed successfully") + + println(notifications.map { it.params.uri }) + // Verify that the notification was sent + await untilAsserted { + assertEquals( + 2, + notifications.size, + "Notification should be sent when resource 1 and resource 2 was deleted", + ) + } + + val deletedResources = listOf(notifications[0].params.uri, notifications[1].params.uri) + assertTrue( + deletedResources.contains(testResourceUri1), + "Notification should contain the removed resource 1 URI", + ) + assertTrue( + deletedResources.contains(testResourceUri2), + "Notification should contain the removed resource 2 URI", + ) + } + + @Test + fun `notification should not be send when removed resource does not exists`() = runTest { + // Track notifications + var resourceListChangedNotificationReceived = false + client.setNotificationHandler( + Method.Defined.NotificationsResourcesListChanged, + ) { + resourceListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Try to remove a non-existent resource + val result = server.removeResource("non-existent-resource") + + // Verify the result + assertFalse(result, "Removing non-existent resource should return false") + assertFalse( + resourceListChangedNotificationReceived, + "No notification should be sent when resource doesn't exist", + ) + } +} diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationTest.kt new file mode 100644 index 00000000..52a42613 --- /dev/null +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesNotificationTest.kt @@ -0,0 +1,150 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceResult +import io.modelcontextprotocol.kotlin.sdk.types.ResourceListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.TextResourceContents +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.test.runTest +import org.awaitility.kotlin.await +import org.awaitility.kotlin.untilAsserted +import org.junit.jupiter.api.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +class ServerResourcesNotificationTest : AbstractServerFeaturesTest() { + + override fun getServerCapabilities(): ServerCapabilities = ServerCapabilities( + resources = ServerCapabilities.Resources(true, null), + ) + + @Test + fun `addResource should send notification`() = runTest { + // Configure notification handler + var resourceListChangedNotificationReceived = false + client.setNotificationHandler( + Method.Defined.NotificationsResourcesListChanged, + ) { + resourceListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Add a resource + val testResourceUri = "test://resource" + server.addResource( + uri = testResourceUri, + name = "Test Resource", + description = "A test resource", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content", + uri = testResourceUri, + mimeType = "text/plain", + ), + ), + ) + } + + // Remove the resource + val result = server.removeResource(testResourceUri) + + // Verify the resource was removed + assertTrue(result, "Resource should be removed successfully") + + // Verify that the notification was sent + await untilAsserted { + assertTrue(resourceListChangedNotificationReceived, "Notification should be sent when resource is added") + } + } + + @Test + fun `removeResources should remove multiple resources and send two notifications`() = runTest { + // Configure notification handler + var resourceListChangedNotificationReceivedCount = 0 + client.setNotificationHandler( + Method.Defined.NotificationsResourcesListChanged, + ) { + resourceListChangedNotificationReceivedCount += 1 + CompletableDeferred(Unit) + } + + // Add resources + val testResourceUri1 = "test://resource1" + val testResourceUri2 = "test://resource2" + + server.addResource( + uri = testResourceUri1, + name = "Test Resource 1", + description = "A test resource 1", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content 1", + uri = testResourceUri1, + mimeType = "text/plain", + ), + ), + ) + } + server.addResource( + uri = testResourceUri2, + name = "Test Resource 2", + description = "A test resource 2", + mimeType = "text/plain", + ) { + ReadResourceResult( + contents = listOf( + TextResourceContents( + text = "Test resource content 2", + uri = testResourceUri2, + mimeType = "text/plain", + ), + ), + ) + } + + // Remove the resources + val result = server.removeResources(listOf(testResourceUri1, testResourceUri2)) + + // Verify the resources were removed + assertEquals(2, result, "Both resources should be removed") + + // Verify that the notifications were sent twice + await untilAsserted { + assertEquals( + 4, + resourceListChangedNotificationReceivedCount, + "Two notifications should be sent when resources are added and two when removed", + ) + } + } + + @Test + fun `notification should not be send when removed resource does not exists`() = runTest { + // Track notifications + var resourceListChangedNotificationReceived = false + client.setNotificationHandler( + Method.Defined.NotificationsResourcesListChanged, + ) { + resourceListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Try to remove a non-existent resource + val result = server.removeResource("non-existent-resource") + + // Verify the result + assertFalse(result, "Removing non-existent resource should return false") + assertFalse( + resourceListChangedNotificationReceived, + "No notification should be sent when resource doesn't exist", + ) + } +} diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesTest.kt index ad73c834..fafc3c86 100644 --- a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesTest.kt +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerResourcesTest.kt @@ -21,7 +21,7 @@ class ServerResourcesTest : AbstractServerFeaturesTest() { ) @Test - fun `removeResource should remove a resource and send notification`() = runTest { + fun `removeResource should remove a resource`() = runTest { // Add a resource val testResourceUri = "test://resource" server.addResource( @@ -49,7 +49,7 @@ class ServerResourcesTest : AbstractServerFeaturesTest() { } @Test - fun `removeResources should remove multiple resources and send notification`() = runTest { + fun `removeResources should remove multiple resources`() = runTest { // Add resources val testResourceUri1 = "test://resource1" val testResourceUri2 = "test://resource2" diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsNotificationTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsNotificationTest.kt new file mode 100644 index 00000000..5892532e --- /dev/null +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsNotificationTest.kt @@ -0,0 +1,100 @@ +package io.modelcontextprotocol.kotlin.sdk.server + +import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult +import io.modelcontextprotocol.kotlin.sdk.types.Method +import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import io.modelcontextprotocol.kotlin.sdk.types.ToolListChangedNotification +import io.modelcontextprotocol.kotlin.sdk.types.ToolSchema +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.test.runTest +import org.awaitility.kotlin.await +import org.awaitility.kotlin.untilAsserted +import org.junit.jupiter.api.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +class ServerToolsNotificationTest : AbstractServerFeaturesTest() { + + override fun getServerCapabilities(): ServerCapabilities = ServerCapabilities( + tools = ServerCapabilities.Tools(true), + ) + + @Test + fun `addTool should send notification`() = runTest { + // Configure notification handler + var toolListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + toolListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Add a tool + server.addTool("test-tool", "Test Tool", ToolSchema()) { + CallToolResult(listOf(TextContent("Test result"))) + } + + // Remove the tool + val result = server.removeTool("test-tool") + + // Verify the tool was removed + assertTrue(result, "Tool should be removed successfully") + + // Verify that the notification was sent + await untilAsserted { + assertTrue(toolListChangedNotificationReceived, "Notification should be sent when tool is added") + } + } + + @Test + fun `removeTools should remove multiple tools and send two notifications`() = runTest { + // Configure notification handler + var toolListChangedNotificationReceivedCount = 0 + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + toolListChangedNotificationReceivedCount += 1 + CompletableDeferred(Unit) + } + + // Add tools + server.addTool("test-tool-1", "Test Tool 1") { + CallToolResult(listOf(TextContent("Test result 1"))) + } + server.addTool("test-tool-2", "Test Tool 2") { + CallToolResult(listOf(TextContent("Test result 2"))) + } + + // Remove the tools + val result = server.removeTools(listOf("test-tool-1", "test-tool-2")) + // Verify the tools were removed + assertEquals(2, result, "Both tools should be removed") + + // Verify that the notifications were sent twice + await untilAsserted { + assertEquals( + 4, + toolListChangedNotificationReceivedCount, + "Two notifications should be sent when tools are added and two when removed", + ) + } + } + + @Test + fun `notification should not be send when removed tool does not exists`() = runTest { + // Track notifications + var toolListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + toolListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Try to remove a non-existent tool + val result = server.removeTool("non-existent-tool") + // Close the server to stop processing further events and flush notifications + server.close() + + // Verify the result + assertFalse(result, "Removing non-existent tool should return false") + assertFalse(toolListChangedNotificationReceived, "No notification should be sent when tool doesn't exist") + } +} diff --git a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsTest.kt b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsTest.kt index eae2eb23..c2e82015 100644 --- a/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsTest.kt +++ b/kotlin-sdk-test/src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/ServerToolsTest.kt @@ -22,7 +22,14 @@ class ServerToolsTest : AbstractServerFeaturesTest() { ) @Test - fun `removeTool should remove a tool`() = runTest { + fun `removeTool should remove a tool and do not send notification`() = runTest { + // Configure notification handler + var toolListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + toolListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + // Add a tool server.addTool("test-tool", "Test Tool", ToolSchema()) { CallToolResult(listOf(TextContent("Test result"))) @@ -33,6 +40,45 @@ class ServerToolsTest : AbstractServerFeaturesTest() { // Verify the tool was removed assertTrue(result, "Tool should be removed successfully") + + // Verify that the notification was not sent + assertFalse( + toolListChangedNotificationReceived, + "No notification should be sent when tools capability is not supported", + ) + } + + @Test + fun `removeTools should remove multiple tools`() = runTest { + // Configure notification handler + var toolListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + toolListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + + // Add tools + server.addTool("test-tool-1", "Test Tool 1") { + CallToolResult(listOf(TextContent("Test result 1"))) + } + server.addTool("test-tool-2", "Test Tool 2") { + CallToolResult(listOf(TextContent("Test result 2"))) + } + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + throw IllegalStateException("Notification should not be sent") + } + + // Remove the tools + val result = server.removeTools(listOf("test-tool-1", "test-tool-2")) + + // Verify the tools were removed + assertEquals(2, result, "Both tools should be removed") + + // Verify that the notification was not sent + assertFalse( + toolListChangedNotificationReceived, + "No notification should be sent when tools capability is not supported", + ) } @Test @@ -54,6 +100,12 @@ class ServerToolsTest : AbstractServerFeaturesTest() { @Test fun `removeTool should throw when tools capability is not supported`() = runTest { + var toolListChangedNotificationReceived = false + client.setNotificationHandler(Method.Defined.NotificationsToolsListChanged) { + toolListChangedNotificationReceived = true + CompletableDeferred(Unit) + } + // Create server without tools capability val serverOptions = ServerOptions( capabilities = ServerCapabilities(), @@ -68,22 +120,11 @@ class ServerToolsTest : AbstractServerFeaturesTest() { server.removeTool("test-tool") } assertEquals("Server does not support tools capability.", exception.message) - } - @Test - fun `removeTools should remove multiple tools`() = runTest { - // Add tools - server.addTool("test-tool-1", "Test Tool 1") { - CallToolResult(listOf(TextContent("Test result 1"))) - } - server.addTool("test-tool-2", "Test Tool 2") { - CallToolResult(listOf(TextContent("Test result 2"))) - } - - // Remove the tools - val result = server.removeTools(listOf("test-tool-1", "test-tool-2")) - - // Verify the tools were removed - assertEquals(2, result, "Both tools should be removed") + // Verify that the notification was not sent + assertFalse( + toolListChangedNotificationReceived, + "No notification should be sent when tools capability is not supported", + ) } }