Skip to content

Commit dc71fc9

Browse files
tiginamariatiginamaria
committed
Add id to session
Support requests proxying from server to session by id
1 parent 07f661a commit dc71fc9

File tree

3 files changed

+191
-6
lines changed

3 files changed

+191
-6
lines changed

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,22 +65,37 @@ 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/CreateElicitationRequest$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/CreateElicitationRequest$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/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/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/Implementation;
78+
public final fun getSession (Ljava/lang/String;)Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;
79+
public final fun getSessionOrThrow (Ljava/lang/String;)Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;
80+
public final fun getSessions ()Ljava/util/Map;
7481
public final fun getTools ()Ljava/util/Map;
82+
public final fun listRoots (Ljava/lang/String;Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
83+
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;
7584
public final fun onClose (Lkotlin/jvm/functions/Function0;)V
7685
public final fun onConnect (Lkotlin/jvm/functions/Function0;)V
7786
public final fun onInitialized (Lkotlin/jvm/functions/Function0;)V
87+
public final fun ping (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
7888
public final fun removePrompt (Ljava/lang/String;)Z
7989
public final fun removePrompts (Ljava/util/List;)I
8090
public final fun removeResource (Ljava/lang/String;)Z
8191
public final fun removeResources (Ljava/util/List;)I
8292
public final fun removeTool (Ljava/lang/String;)Z
8393
public final fun removeTools (Ljava/util/List;)I
94+
public final fun sendLoggingMessage (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/LoggingMessageNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
95+
public final fun sendPromptListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
96+
public final fun sendResourceListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
97+
public final fun sendResourceUpdated (Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/ResourceUpdatedNotification;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
98+
public final fun sendToolListChanged (Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
8499
}
85100

86101
public final class io/modelcontextprotocol/kotlin/sdk/server/ServerOptions : io/modelcontextprotocol/kotlin/sdk/shared/ProtocolOptions {
@@ -98,10 +113,13 @@ public class io/modelcontextprotocol/kotlin/sdk/server/ServerSession : io/modelc
98113
public static synthetic fun createElicitation$default (Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;Ljava/lang/String;Lio/modelcontextprotocol/kotlin/sdk/CreateElicitationRequest$RequestedSchema;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
99114
public final fun createMessage (Lio/modelcontextprotocol/kotlin/sdk/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
100115
public static synthetic fun createMessage$default (Lio/modelcontextprotocol/kotlin/sdk/server/ServerSession;Lio/modelcontextprotocol/kotlin/sdk/CreateMessageRequest;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
116+
public fun equals (Ljava/lang/Object;)Z
101117
public final fun getClientCapabilities ()Lio/modelcontextprotocol/kotlin/sdk/ClientCapabilities;
102118
public final fun getClientVersion ()Lio/modelcontextprotocol/kotlin/sdk/Implementation;
103119
protected final fun getInstructions ()Ljava/lang/String;
104120
protected final fun getServerInfo ()Lio/modelcontextprotocol/kotlin/sdk/Implementation;
121+
public final fun getSessionId ()Ljava/lang/String;
122+
public fun hashCode ()I
105123
public final fun listRoots (Lkotlinx/serialization/json/JsonObject;Lio/modelcontextprotocol/kotlin/sdk/shared/RequestOptions;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
106124
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;
107125
public fun onClose ()V

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

Lines changed: 155 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,12 @@ package io.modelcontextprotocol.kotlin.sdk.server
33
import io.github.oshai.kotlinlogging.KotlinLogging
44
import io.modelcontextprotocol.kotlin.sdk.CallToolRequest
55
import io.modelcontextprotocol.kotlin.sdk.CallToolResult
6+
import io.modelcontextprotocol.kotlin.sdk.CreateElicitationRequest.RequestedSchema
7+
import io.modelcontextprotocol.kotlin.sdk.CreateElicitationResult
8+
import io.modelcontextprotocol.kotlin.sdk.CreateMessageRequest
9+
import io.modelcontextprotocol.kotlin.sdk.CreateMessageResult
610
import io.modelcontextprotocol.kotlin.sdk.EmptyJsonObject
11+
import io.modelcontextprotocol.kotlin.sdk.EmptyRequestResult
712
import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
813
import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
914
import io.modelcontextprotocol.kotlin.sdk.Implementation
@@ -13,23 +18,27 @@ import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesRequest
1318
import io.modelcontextprotocol.kotlin.sdk.ListResourceTemplatesResult
1419
import io.modelcontextprotocol.kotlin.sdk.ListResourcesRequest
1520
import io.modelcontextprotocol.kotlin.sdk.ListResourcesResult
21+
import io.modelcontextprotocol.kotlin.sdk.ListRootsResult
1622
import io.modelcontextprotocol.kotlin.sdk.ListToolsRequest
1723
import io.modelcontextprotocol.kotlin.sdk.ListToolsResult
24+
import io.modelcontextprotocol.kotlin.sdk.LoggingMessageNotification
1825
import io.modelcontextprotocol.kotlin.sdk.Method
1926
import io.modelcontextprotocol.kotlin.sdk.Prompt
2027
import io.modelcontextprotocol.kotlin.sdk.PromptArgument
2128
import io.modelcontextprotocol.kotlin.sdk.ReadResourceRequest
2229
import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
2330
import io.modelcontextprotocol.kotlin.sdk.Resource
31+
import io.modelcontextprotocol.kotlin.sdk.ResourceUpdatedNotification
2432
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
2533
import io.modelcontextprotocol.kotlin.sdk.TextContent
2634
import io.modelcontextprotocol.kotlin.sdk.Tool
2735
import io.modelcontextprotocol.kotlin.sdk.ToolAnnotations
2836
import io.modelcontextprotocol.kotlin.sdk.shared.ProtocolOptions
37+
import io.modelcontextprotocol.kotlin.sdk.shared.RequestOptions
2938
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
3039
import kotlinx.atomicfu.atomic
3140
import kotlinx.atomicfu.update
32-
import kotlinx.collections.immutable.persistentListOf
41+
import kotlinx.collections.immutable.persistentMapOf
3342
import kotlinx.coroutines.CancellationException
3443
import kotlinx.serialization.json.JsonObject
3544

@@ -45,7 +54,7 @@ public class ServerOptions(public val capabilities: ServerCapabilities, enforceS
4554
ProtocolOptions(enforceStrictCapabilities = enforceStrictCapabilities)
4655

4756
/**
48-
* An MCP server on top of a pluggable transport.
57+
* An MCP server is responsible for storing features and handling new connections.
4958
*
5059
* This server automatically responds to the initialization flow as initiated by the client.
5160
* You can register tools, prompts, and resources using [addTool], [addPrompt], and [addResource].
@@ -79,7 +88,24 @@ public open class Server(
7988
block: Server.() -> Unit = {},
8089
) : this(serverInfo, options, { instructions }, block)
8190

82-
private val sessions = atomic(persistentListOf<ServerSession>())
91+
private val sessionRegistry = atomic(persistentMapOf<String, ServerSession>())
92+
93+
/**
94+
* Returns a read-only view of the current server sessions.
95+
*/
96+
public val sessions: Map<String, ServerSession>
97+
get() = sessionRegistry.value
98+
99+
/**
100+
* Gets a server session by its ID.
101+
*/
102+
public fun getSession(sessionId: String): ServerSession? = sessions[sessionId]
103+
104+
/**
105+
* Gets a server session by its ID or throws an exception if the session doesn't exist.
106+
*/
107+
public fun getSessionOrThrow(sessionId: String): ServerSession =
108+
sessions[sessionId] ?: throw IllegalArgumentException("Session not found: $sessionId")
83109

84110
@Suppress("ktlint:standard:backing-property-naming")
85111
private var _onInitialized: (() -> Unit) = {}
@@ -107,7 +133,10 @@ public open class Server(
107133

108134
public suspend fun close() {
109135
logger.debug { "Closing MCP server" }
110-
sessions.value.forEach { session -> session.close() }
136+
sessions.forEach { (sessionId, session) ->
137+
logger.info { "Closing session $sessionId" }
138+
session.close()
139+
}
111140
_onClose()
112141
}
113142

@@ -171,12 +200,12 @@ public open class Server(
171200
// Register cleanup handler to remove session from list when it closes
172201
session.onClose {
173202
logger.debug { "Removing closed session from active sessions list" }
174-
sessions.update { list -> list.remove(session) }
203+
sessionRegistry.update { sessions -> sessions.remove(session.sessionId) }
175204
}
176205
logger.debug { "Server session connecting to transport" }
177206
session.connect(transport)
178207
logger.debug { "Server session successfully connected to transport" }
179-
sessions.update { sessions -> sessions.add(session) }
208+
sessionRegistry.update { sessions -> sessions.put(session.sessionId, session) }
180209

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

‎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
@@ -32,14 +32,24 @@ import kotlinx.atomicfu.AtomicRef
3232
import kotlinx.atomicfu.atomic
3333
import kotlinx.coroutines.CompletableDeferred
3434
import kotlinx.serialization.json.JsonObject
35+
import kotlin.uuid.ExperimentalUuidApi
36+
import kotlin.uuid.Uuid
3537

3638
private val logger = KotlinLogging.logger {}
3739

40+
/**
41+
* Represents a server session.
42+
*/
43+
@Suppress("TooManyFunctions")
3844
public open class ServerSession(
3945
protected val serverInfo: Implementation,
4046
options: ServerOptions,
4147
protected val instructions: String?,
4248
) : Protocol(options) {
49+
50+
@OptIn(ExperimentalUuidApi::class)
51+
public val sessionId: String = Uuid.random().toString()
52+
4353
@Suppress("ktlint:standard:backing-property-naming")
4454
private var _onInitialized: (() -> Unit) = {}
4555

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

0 commit comments

Comments
 (0)