Server feature list changed and resource updated notifications (#441)
<!-- Provide a brief summary of your changes --> ## Motivation and Context <!-- Why is this change needed? What problem does it solve? --> Implemented thread-safe notification service following eventual consistency principle. Fixes https://github.com/modelcontextprotocol/kotlin-sdk/issues/249 Contracts: * If the client connection and feature update happened in the same thread, the timestamps of the events will be consequent ``` server.connect(client) server.add(feature) // client will receive the notification ``` ``` server.add(feature) server.connect(client) // client will NOT receive the notification ``` * If the client connection and feature update happened in the different threads, both options can be possible The flow of the notification is explained in the diagram: <img width="663" height="413" alt="image" src="https://github.com/user-attachments/assets/e2c4c108-53dd-408d-97ef-97110ce5c979" /> Documentation references: https://modelcontextprotocol.io/specification/2025-06-18/server/prompts#list-changed-notification https://modelcontextprotocol.io/specification/2025-06-18/server/tools#list-changed-notification https://modelcontextprotocol.io/specification/2025-06-18/server/resources#list-changed-notification https://modelcontextprotocol.io/specification/2025-06-18/server/resources#subscriptions Questions: * I haven't found the subscription confirmed response message content docs, how it should look like? ## How Has This Been Tested? <!-- Have you tested this in a real application? Which scenarios were tested? --> Multiple tests with various notification scenarios were added ## Breaking Changes <!-- Will users need to update their code or configurations? --> ## Types of changes <!-- What types of changes does your code introduce? Put an `x` in all the boxes that apply: --> - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update ## Checklist <!-- Go over all the following points, and put an `x` in all the boxes that apply. --> - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [x] New and existing tests pass locally - [x] I have added appropriate error handling - [x] I have added or updated documentation as needed ## Additional context <!-- Add any other context, implementation notes, or design decisions -->
This commit is contained in:
Maria Tigina
GitHub
parent
cc9b50ff76
commit
4645ebfd17
5 changed files with 1307 additions and 60 deletions
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,56 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.server
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CallToolRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.GetPromptRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.GetPromptResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Prompt
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Resource
|
||||
import io.modelcontextprotocol.kotlin.sdk.types.Tool
|
||||
|
||||
internal typealias FeatureKey = String
|
||||
|
||||
/**
|
||||
* Represents a feature with an associated unique key.
|
||||
*/
|
||||
internal interface Feature {
|
||||
val key: FeatureKey
|
||||
}
|
||||
|
||||
/**
|
||||
* A wrapper class representing a registered tool on the server.
|
||||
*
|
||||
* @property tool The tool definition.
|
||||
* @property handler A suspend function to handle the tool call requests.
|
||||
*/
|
||||
public data class RegisteredTool(val tool: Tool, val handler: suspend (CallToolRequest) -> CallToolResult) : Feature {
|
||||
override val key: String = tool.name
|
||||
}
|
||||
|
||||
/**
|
||||
* A wrapper class representing a registered prompt on the server.
|
||||
*
|
||||
* @property prompt The prompt definition.
|
||||
* @property messageProvider A suspend function that returns the prompt content when requested by the client.
|
||||
*/
|
||||
public data class RegisteredPrompt(
|
||||
val prompt: Prompt,
|
||||
val messageProvider: suspend (GetPromptRequest) -> GetPromptResult,
|
||||
) : Feature {
|
||||
override val key: String = prompt.name
|
||||
}
|
||||
|
||||
/**
|
||||
* A wrapper class representing a registered resource on the server.
|
||||
*
|
||||
* @property resource The resource definition.
|
||||
* @property readHandler A suspend function to handle read requests for this resource.
|
||||
*/
|
||||
public data class RegisteredResource(
|
||||
val resource: Resource,
|
||||
val readHandler: suspend (ReadResourceRequest) -> ReadResourceResult,
|
||||
) : Feature {
|
||||
override val key: String = resource.uri
|
||||
}
|
||||
|
|
@ -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<FeatureKey, Long>())
|
||||
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<NotificationEvent>,
|
||||
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<NotificationEvent>(
|
||||
extraBufferCapacity = notificationBufferCapacity,
|
||||
onBufferOverflow = BufferOverflow.SUSPEND,
|
||||
)
|
||||
|
||||
/** Active emit jobs. */
|
||||
private val emitJobs = atomic(persistentSetOf<Job>())
|
||||
|
||||
/** Notification jobs associated with sessions. */
|
||||
private val sessionNotificationJobs = atomic(persistentMapOf<ServerSessionKey, SessionNotificationJob>())
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
|
@ -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<T : Feature>(private val featureType: String) {
|
|||
internal val values: Map<FeatureKey, T>
|
||||
get() = registry.value
|
||||
|
||||
private val listeners = atomic(persistentListOf<FeatureListener>())
|
||||
|
||||
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<T : Feature>(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<T : Feature>(private val featureType: String) {
|
|||
*/
|
||||
internal fun addAll(features: List<T>) {
|
||||
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<T : Feature>(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<T : Feature>(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<T : Feature>(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) }
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<ServerSessionKey, ServerSession>
|
||||
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<RegisteredTool>("Tool")
|
||||
private val promptRegistry = FeatureRegistry<RegisteredPrompt>("Prompt")
|
||||
private val resourceRegistry = FeatureRegistry<RegisteredResource>("Resource")
|
||||
@OptIn(ExperimentalTime::class)
|
||||
private val notificationService = FeatureNotificationService()
|
||||
|
||||
private val sessionRegistry = ServerSessionRegistry()
|
||||
|
||||
private val toolRegistry = FeatureRegistry<RegisteredTool>("Tool").apply {
|
||||
if (options.capabilities.tools?.listChanged ?: false) {
|
||||
addListener(notificationService.toolListChangedListener)
|
||||
}
|
||||
}
|
||||
private val promptRegistry = FeatureRegistry<RegisteredPrompt>("Prompt").apply {
|
||||
if (options.capabilities.prompts?.listChanged ?: false) {
|
||||
addListener(notificationService.promptListChangedListener)
|
||||
}
|
||||
}
|
||||
private val resourceRegistry = FeatureRegistry<RegisteredResource>("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<ServerSessionKey, ServerSession>
|
||||
get() = sessionRegistry.sessions
|
||||
|
||||
/**
|
||||
* Provides a snapshot of all tools currently registered in the server
|
||||
*/
|
||||
public val tools: Map<String, RegisteredTool>
|
||||
get() = toolRegistry.values
|
||||
|
||||
/**
|
||||
* Provides a snapshot of all prompts currently registered in the server
|
||||
*/
|
||||
public val prompts: Map<String, RegisteredPrompt>
|
||||
get() = promptRegistry.values
|
||||
|
||||
/**
|
||||
* Provides a snapshot of all resources currently registered in the server
|
||||
*/
|
||||
public val resources: Map<String, RegisteredResource>
|
||||
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<ListResourceTemplatesRequest>(Method.Defined.ResourcesTemplatesList) { _, _ ->
|
||||
handleListResourceTemplates()
|
||||
}
|
||||
if (options.capabilities.resources?.subscribe ?: false) {
|
||||
session.setRequestHandler<SubscribeRequest>(Method.Defined.ResourcesSubscribe) { request, _ ->
|
||||
handleSubscribeResources(session, request)
|
||||
// Does not return any confirmation as the structure is not stated in the protocol
|
||||
null
|
||||
}
|
||||
session.setRequestHandler<UnsubscribeRequest>(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 <T : Notification> setNotificationHandler(method: Method, handler: (notification: T) -> Deferred<Unit>) {
|
||||
sessions.forEach { (_, session) ->
|
||||
session.setNotificationHandler(method, handler)
|
||||
}
|
||||
}
|
||||
|
||||
public fun removeNotificationHandler(method: Method) {
|
||||
sessions.forEach { (_, session) ->
|
||||
session.removeNotificationHandler(method)
|
||||
}
|
||||
}
|
||||
|
||||
public fun <T : Notification> setNotificationHandler(
|
||||
sessionId: String,
|
||||
method: Method,
|
||||
handler: (notification: T) -> Deferred<Unit>,
|
||||
) {
|
||||
sessionRegistry.getSessionOrNull(sessionId)?.setNotificationHandler(method, handler)
|
||||
}
|
||||
|
||||
public fun removeNotificationHandler(sessionId: String, method: Method) {
|
||||
sessionRegistry.getSessionOrNull(sessionId)?.removeNotificationHandler(method)
|
||||
}
|
||||
// End the notification handling section
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue