Skip to content

Commit 0e6a41d

Browse files
tiginamariatiginamaria
authored
Add server session id (#381)
<!-- Provide a brief summary of your changes --> ## Motivation and Context <!-- Why is this change needed? What problem does it solve? --> Add ids for server sessions (auto generated when created) Support methods proxying from server to server session by id ## How Has This Been Tested? <!-- Have you tested this in a real application? Which scenarios were tested? --> ## 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. --> - [ ] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [ ] My code follows the repository's style guidelines - [ ] New and existing tests pass locally - [ ] I have added appropriate error handling - [ ] I have added or updated documentation as needed ## Additional context <!-- Add any other context, implementation notes, or design decisions -->
1 parent 21393e7 commit 0e6a41d

File tree

4 files changed

+239
-8
lines changed

4 files changed

+239
-8
lines changed

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,22 +65,35 @@ public class io/modelcontextprotocol/kotlin/sdk/server/Server {
6565
public final fun addTools (Ljava/util/List;)V
6666
public final fun close (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
6767
public final fun connect (Lio/modelcontextprotocol/kotlin/sdk/shared/Transport;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
68+
public final fun createElicitation (Ljava/lang/String;Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/ElicitRequestParams$RequestedSchema;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
69+
public static synthetic fun createElicitation$default (Lio/modelcontextprotocol/kotlin/sdk/server/Server;Ljava/lang/String;Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/ElicitRequestParams$RequestedSchema;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
70+
public final fun createMessage (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
71+
public static synthetic fun createMessage$default (Lio/modelcontextprotocol/kotlin/sdk/server/Server;Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
6872
public final fun createSession (Lio/modelcontextprotocol/kotlin/sdk/shared/Transport;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
6973
protected final fun getInstructionsProvider ()Lkotlin/jvm/functions/Function0;
7074
protected final fun getOptions ()Lio/modelcontextprotocol/kotlin/sdk/server/ServerOptions;
7175
public final fun getPrompts ()Ljava/util/Map;
7276
public final fun getResources ()Ljava/util/Map;
7377
protected final fun getServerInfo ()Lio/modelcontextprotocol/kotlin/sdk/types/Implementation;
78+
public final fun getSessions ()Ljava/util/Map;
7479
public final fun getTools ()Ljava/util/Map;
80+
public final fun listRoots (Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
81+
public static synthetic fun listRoots$default (Lio/modelcontextprotocol/kotlin/sdk/server/Server;Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
7582
public final fun onClose (Lkotlin/jvm/functions/Function0;)V
7683
public final fun onConnect (Lkotlin/jvm/functions/Function0;)V
7784
public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V
85+
public final fun ping (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
7886
public final fun removePrompt (Ljava/lang/String;)Z
7987
public final fun removePrompts (Ljava/util/List;)I
8088
public final fun removeResource (Ljava/lang/String;)Z
8189
public final fun removeResources (Ljava/util/List;)I
8290
public final fun removeTool (Ljava/lang/String;)Z
8391
public final fun removeTools (Ljava/util/List;)I
92+
public final fun sendLoggingMessage (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/LoggingMessageNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
93+
public final fun sendPromptListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
94+
public final fun sendResourceListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
95+
public final fun sendResourceUpdated (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/ResourceUpdatedNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
96+
public final fun sendToolListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
8497
}
8598

8699
public final class io/modelcontextprotocol/kotlin/sdk/server/ServerOptions : io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions {
@@ -98,10 +111,13 @@ public class io/modelcontextprotocol/kotlin/sdk/server/ServerSession : io/modelc
98111
public static synthetic fun createElicitation$default (Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/types/ElicitRequestParams$RequestedSchema;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
99112
public final fun createMessage (Lio/modelcontextprotocol/kotlin/sdk/types/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
100113
public static synthetic fun createMessage$default (Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;Lio/modelcontextprotocol/kotlin/sdk/types/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
114+
public fun equals (Ljava/lang/Object;)Z
101115
public final fun getClientCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/types/ClientCapabilities;
102116
public final fun getClientVersion ()Lio/modelcontextprotocol/kotlin/sdk/types/Implementation;
103117
protected final fun getInstructions ()Ljava/lang/String;
104118
protected final fun getServerInfo ()Lio/modelcontextprotocol/kotlin/sdk/types/Implementation;
119+
public final fun getSessionId ()Ljava/lang/String;
120+
public fun hashCode ()I
105121
public final fun listRoots (Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
106122
public static synthetic fun listRoots$default (Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
107123
public fun onClose ()V

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

Lines changed: 145 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,16 @@ package io.modelcontextprotocol.kotlin.sdk.server
22

33
import io.github.oshai.kotlinlogging.KotlinLogging
44
import io.modelcontextprotocol.kotlin.sdk.shared.ProtocolOptions
5+
import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
56
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
67
import io.modelcontextprotocol.kotlin.sdk.types.CallToolRequest
78
import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult
9+
import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageRequest
10+
import io.modelcontextprotocol.kotlin.sdk.types.CreateMessageResult
11+
import io.modelcontextprotocol.kotlin.sdk.types.ElicitRequestParams
12+
import io.modelcontextprotocol.kotlin.sdk.types.ElicitResult
13+
import io.modelcontextprotocol.kotlin.sdk.types.EmptyJsonObject
14+
import io.modelcontextprotocol.kotlin.sdk.types.EmptyResult
815
import io.modelcontextprotocol.kotlin.sdk.types.GetPromptRequest
916
import io.modelcontextprotocol.kotlin.sdk.types.GetPromptResult
1017
import io.modelcontextprotocol.kotlin.sdk.types.Implementation
@@ -14,22 +21,22 @@ import io.modelcontextprotocol.kotlin.sdk.types.ListResourceTemplatesRequest
1421
import io.modelcontextprotocol.kotlin.sdk.types.ListResourceTemplatesResult
1522
import io.modelcontextprotocol.kotlin.sdk.types.ListResourcesRequest
1623
import io.modelcontextprotocol.kotlin.sdk.types.ListResourcesResult
24+
import io.modelcontextprotocol.kotlin.sdk.types.ListRootsResult
1725
import io.modelcontextprotocol.kotlin.sdk.types.ListToolsRequest
1826
import io.modelcontextprotocol.kotlin.sdk.types.ListToolsResult
27+
import io.modelcontextprotocol.kotlin.sdk.types.LoggingMessageNotification
1928
import io.modelcontextprotocol.kotlin.sdk.types.Method
2029
import io.modelcontextprotocol.kotlin.sdk.types.Prompt
2130
import io.modelcontextprotocol.kotlin.sdk.types.PromptArgument
2231
import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceRequest
2332
import io.modelcontextprotocol.kotlin.sdk.types.ReadResourceResult
2433
import io.modelcontextprotocol.kotlin.sdk.types.Resource
34+
import io.modelcontextprotocol.kotlin.sdk.types.ResourceUpdatedNotification
2535
import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities
2636
import io.modelcontextprotocol.kotlin.sdk.types.TextContent
2737
import io.modelcontextprotocol.kotlin.sdk.types.Tool
2838
import io.modelcontextprotocol.kotlin.sdk.types.ToolAnnotations
2939
import io.modelcontextprotocol.kotlin.sdk.types.ToolSchema
30-
import kotlinx.atomicfu.atomic
31-
import kotlinx.atomicfu.update
32-
import kotlinx.collections.immutable.persistentListOf
3340
import kotlinx.coroutines.CancellationException
3441
import kotlinx.serialization.json.JsonObject
3542

@@ -45,7 +52,7 @@ public class ServerOptions(public val capabilities: ServerCapabilities, enforceS
4552
ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities)
4653

4754
/**
48-
* An MCP server on top of a pluggable transport.
55+
* An MCP server is responsible for storing features and handling new connections.
4956
*
5057
* This server automatically responds to the initialization flow as initiated by the client.
5158
* You can register tools, prompts, and resources using [addTool], [addPrompt], and [addResource].
@@ -79,7 +86,13 @@ public open class Server(
7986
block: Server.() -> Unit = {},
8087
) : this(serverInfo, options, { instructions }, block)
8188

82-
private val sessions = atomic(persistentListOf<ServerSession>())
89+
private val sessionRegistry = ServerSessionRegistry()
90+
91+
/**
92+
* Provides a snapshot of all sessions currently registered in the server
93+
*/
94+
public val sessions: Map<ServerSessionKey, ServerSession>
95+
get() = sessionRegistry.sessions
8396

8497
@Suppress("ktlint:standard:backing-property-naming")
8598
private var _onInitialized: (() -> Unit) = {}
@@ -107,7 +120,10 @@ public open class Server(
107120

108121
public suspend fun close() {
109122
logger.debug { "Closing MCP server" }
110-
sessions.value.forEach { session -> session.close() }
123+
sessions.forEach { (sessionId, session) ->
124+
logger.info { "Closing session $sessionId" }
125+
session.close()
126+
}
111127
_onClose()
112128
}
113129

@@ -171,12 +187,12 @@ public open class Server(
171187
// Register cleanup handler to remove session from list when it closes
172188
session.onClose {
173189
logger.debug { "Removing closed session from active sessions list" }
174-
sessions.update { list -> list.remove(session) }
190+
sessionRegistry.removeSession(session.sessionId)
175191
}
176192
logger.debug { "Server session connecting to transport" }
177193
session.connect(transport)
178194
logger.debug { "Server session successfully connected to transport" }
179-
sessions.update { sessions -> sessions.add(session) }
195+
sessionRegistry.addSession(session)
180196

181197
_onConnect()
182198
return session
@@ -538,4 +554,125 @@ public open class Server(
538554
// If you have resource templates, return them here. For now, return empty.
539555
return ListResourceTemplatesResult(listOf())
540556
}
557+
558+
// Start the ServerSession redirection section
559+
560+
/**
561+
* Triggers [ServerSession.ping] request for session by provided [sessionId].
562+
* @param sessionId The session ID to ping
563+
*/
564+
public suspend fun ping(sessionId: String): EmptyResult = with(sessionRegistry.getSession(sessionId)) {
565+
ping()
566+
}
567+
568+
/**
569+
* Triggers [ServerSession.createMessage] request for session by provided [sessionId].
570+
*
571+
* @param sessionId The session ID to create a message.
572+
* @param params The parameters for creating a message.
573+
* @param options Optional request options.
574+
* @return The created message result.
575+
* @throws IllegalStateException If the server does not support sampling or if the request fails.
576+
*/
577+
public suspend fun createMessage(
578+
sessionId: String,
579+
params: CreateMessageRequest,
580+
options: RequestOptions? = null,
581+
): CreateMessageResult = with(sessionRegistry.getSession(sessionId)) {
582+
request(params, options)
583+
}
584+
585+
/**
586+
* Triggers [ServerSession.listRoots] request for session by provided [sessionId].
587+
*
588+
* @param sessionId The session ID to list roots for.
589+
* @param params JSON parameters for the request, usually empty.
590+
* @param options Optional request options.
591+
* @return The list of roots.
592+
* @throws IllegalStateException If the server or client does not support roots.
593+
*/
594+
public suspend fun listRoots(
595+
sessionId: String,
596+
params: JsonObject = EmptyJsonObject,
597+
options: RequestOptions? = null,
598+
): ListRootsResult = with(sessionRegistry.getSession(sessionId)) {
599+
listRoots(params, options)
600+
}
601+
602+
/**
603+
* Triggers [ServerSession.createElicitation] request for session by provided [sessionId].
604+
*
605+
* @param sessionId The session ID to create elicitation for.
606+
* @param message The elicitation message.
607+
* @param requestedSchema The requested schema for the elicitation.
608+
* @param options Optional request options.
609+
* @return The created elicitation result.
610+
* @throws IllegalStateException If the server does not support elicitation or if the request fails.
611+
*/
612+
public suspend fun createElicitation(
613+
sessionId: String,
614+
message: String,
615+
requestedSchema: ElicitRequestParams.RequestedSchema,
616+
options: RequestOptions? = null,
617+
): ElicitResult = with(sessionRegistry.getSession(sessionId)) {
618+
createElicitation(message, requestedSchema, options)
619+
}
620+
621+
/**
622+
* Triggers [ServerSession.sendLoggingMessage] for session by provided [sessionId].
623+
*
624+
* @param sessionId The session ID to send the logging message to.
625+
* @param notification The logging message notification.
626+
*/
627+
public suspend fun sendLoggingMessage(sessionId: String, notification: LoggingMessageNotification) {
628+
with(sessionRegistry.getSession(sessionId)) {
629+
sendLoggingMessage(notification)
630+
}
631+
}
632+
633+
/**
634+
* Triggers [ServerSession.sendResourceUpdated] for session by provided [sessionId].
635+
*
636+
* @param sessionId The session ID to send the resource updated notification to.
637+
* @param notification Details of the updated resource.
638+
*/
639+
public suspend fun sendResourceUpdated(sessionId: String, notification: ResourceUpdatedNotification) {
640+
with(sessionRegistry.getSession(sessionId)) {
641+
sendResourceUpdated(notification)
642+
}
643+
}
644+
645+
/**
646+
* Triggers [ServerSession.sendResourceListChanged] for session by provided [sessionId].
647+
*
648+
* @param sessionId The session ID to send the resource list changed notification to.
649+
*/
650+
public suspend fun sendResourceListChanged(sessionId: String) {
651+
with(sessionRegistry.getSession(sessionId)) {
652+
sendResourceListChanged()
653+
}
654+
}
655+
656+
/**
657+
* Triggers [ServerSession.sendToolListChanged] for session by provided [sessionId].
658+
*
659+
* @param sessionId The session ID to send the tool list changed notification to.
660+
*/
661+
public suspend fun sendToolListChanged(sessionId: String) {
662+
with(sessionRegistry.getSession(sessionId)) {
663+
sendToolListChanged()
664+
}
665+
}
666+
667+
/**
668+
* Triggers [ServerSession.sendPromptListChanged] for session by provided [sessionId].
669+
*
670+
* @param sessionId The session ID to send the prompt list changed notification to.
671+
*/
672+
public suspend fun sendPromptListChanged(sessionId: String) {
673+
with(sessionRegistry.getSession(sessionId)) {
674+
sendPromptListChanged()
675+
}
676+
}
677+
// End the ServerSession redirection section
541678
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,14 +35,24 @@ import kotlinx.atomicfu.AtomicRef
3535
import kotlinx.atomicfu.atomic
3636
import kotlinx.coroutines.CompletableDeferred
3737
import kotlinx.serialization.json.JsonObject
38+
import kotlin.uuid.ExperimentalUuidApi
39+
import kotlin.uuid.Uuid
3840

3941
private val logger = KotlinLogging.logger {}
4042

43+
/**
44+
* Represents a server session.
45+
*/
46+
@Suppress("TooManyFunctions")
4147
public open class ServerSession(
4248
protected val serverInfo: Implementation,
4349
options: ServerOptions,
4450
protected val instructions: String?,
4551
) : Protocol(options) {
52+
53+
@OptIn(ExperimentalUuidApi::class)
54+
public val sessionId: String = Uuid.random().toString()
55+
4656
@Suppress("ktlint:standard:backing-property-naming")
4757
private var _onInitialized: (() -> Unit) = {}
4858

@@ -430,4 +440,12 @@ public open class ServerSession(
430440
* @return true if the message should be accepted (not filtered out), false otherwise.
431441
*/
432442
private fun isMessageAccepted(level: LoggingLevel): Boolean = !isMessageIgnored(level)
443+
444+
override fun equals(other: Any?): Boolean {
445+
if (this === other) return true
446+
if (other !is ServerSession) return false
447+
return sessionId == other.sessionId
448+
}
449+
450+
override fun hashCode(): Int = sessionId.hashCode()
433451
}

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
package io.modelcontextprotocol.kotlin.sdk.server
2+
3+
import io.github.oshai.kotlinlogging.KotlinLogging
4+
import kotlinx.atomicfu.atomic
5+
import kotlinx.atomicfu.update
6+
import kotlinx.collections.immutable.persistentMapOf
7+
8+
internal typealias ServerSessionKey = String
9+
10+
/**
11+
* Represents a registry for managing server sessions.
12+
*/
13+
internal class ServerSessionRegistry {
14+
15+
private val logger = KotlinLogging.logger {}
16+
17+
/**
18+
* Atomic variable used to maintain a thread-safe registry of sessions.
19+
* Stores a persistent map where each session is identified by its unique key.
20+
*/
21+
private val registry = atomic(persistentMapOf<String, ServerSession>())
22+
23+
/**
24+
* Returns a read-only view of the current server sessions.
25+
*/
26+
internal val sessions: Map<ServerSessionKey, ServerSession>
27+
get() = registry.value
28+
29+
/**
30+
* Returns a server session by its ID.
31+
* @param sessionId The ID of the session to retrieve.
32+
* @throws IllegalArgumentException If the session doesn't exist.
33+
*/
34+
internal fun getSession(sessionId: ServerSessionKey): ServerSession =
35+
sessions[sessionId] ?: throw IllegalArgumentException("Session not found: $sessionId")
36+
37+
/**
38+
* Returns a server session by its ID, or null if it doesn't exist.
39+
* @param sessionId The ID of the session to retrieve.
40+
*/
41+
internal fun getSessionOrNull(sessionId: ServerSessionKey): ServerSession? = sessions[sessionId]
42+
43+
/**
44+
* Registers a server session.
45+
* @param session The session to register.
46+
*/
47+
internal fun addSession(session: ServerSession) {
48+
logger.info { "Adding session: ${session.sessionId}" }
49+
registry.update { sessions -> sessions.put(session.sessionId, session) }
50+
}
51+
52+
/**
53+
* Removes a server session by its ID.
54+
* @param sessionId The ID of the session to remove.
55+
*/
56+
internal fun removeSession(sessionId: ServerSessionKey) {
57+
logger.info { "Removing session: $sessionId" }
58+
registry.update { sessions -> sessions.remove(sessionId) }
59+
}
60+
}

0 commit comments

Comments
 (0)