Introduce server sessions (#198)
Introduce a `ServerSession` to manage multiple connections to the mcp server. `Protocall` logic of the `Server` was moved to `ServerSession`, while MCP resources logic remained in `Server` ## Motivation and Context 1. Fix the multiple connection issues reported for SSE transport 2. Avoid server duplications during configurations: currently every connection creates a session, but the server stays single ## How Has This Been Tested? Added `SseIntegrationTest` and `WebSocketIntegrationTest` integration tests to check connection with single and multiple clients ## Breaking Changes This change caused several depreciations: 1. The `Server` -> `Protocol` inheritance (which caused single connection issue) should be removed in future releases, now, all the methods inherited from `Protocall` marked as `Deprecated` with a `Warning` 2. Websocket Ktor server initialization changes in order not to create server instance per connection ## Types of changes - [ ] 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 --> --------- Co-authored-by: devcrocod <devcrocod@gmail.com>
This commit is contained in:
Maria Tigina
GitHub
parent
0ed92f5b03
commit
d86813636e
6 changed files with 1191 additions and 602 deletions
|
|
@ -30,6 +30,7 @@ import io.modelcontextprotocol.kotlin.sdk.TextContent
|
|||
import io.modelcontextprotocol.kotlin.sdk.Tool
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.Server
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.ServerSession
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.InMemoryTransport
|
||||
import kotlinx.coroutines.CompletableDeferred
|
||||
|
|
@ -241,7 +242,28 @@ class ClientTest {
|
|||
serverOptions,
|
||||
)
|
||||
|
||||
server.setRequestHandler<InitializeRequest>(Method.Defined.Initialize) { _, _ ->
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(sampling = EmptyJsonObject),
|
||||
),
|
||||
)
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
serverSession.setRequestHandler<InitializeRequest>(Method.Defined.Initialize) { _, _ ->
|
||||
InitializeResult(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ServerCapabilities(
|
||||
|
|
@ -252,32 +274,13 @@ class ClientTest {
|
|||
)
|
||||
}
|
||||
|
||||
server.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
serverSession.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
ListResourcesResult(resources = emptyList(), nextCursor = null)
|
||||
}
|
||||
|
||||
server.setRequestHandler<ListToolsRequest>(Method.Defined.ToolsList) { _, _ ->
|
||||
serverSession.setRequestHandler<ListToolsRequest>(Method.Defined.ToolsList) { _, _ ->
|
||||
ListToolsResult(tools = emptyList(), nextCursor = null)
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(
|
||||
capabilities = ClientCapabilities(sampling = EmptyJsonObject),
|
||||
),
|
||||
)
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
// Server supports resources and tools, but not prompts
|
||||
val caps = client.serverCapabilities
|
||||
assertEquals(ServerCapabilities.Resources(null, null), caps?.resources)
|
||||
|
|
@ -368,24 +371,27 @@ class ClientTest {
|
|||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
// These should not throw
|
||||
val jsonObject = buildJsonObject {
|
||||
put("name", "John")
|
||||
put("age", 30)
|
||||
put("isStudent", false)
|
||||
}
|
||||
server.sendLoggingMessage(
|
||||
serverSession.sendLoggingMessage(
|
||||
LoggingMessageNotification(
|
||||
params = LoggingMessageNotification.Params(
|
||||
level = LoggingLevel.info,
|
||||
|
|
@ -393,11 +399,11 @@ class ClientTest {
|
|||
),
|
||||
),
|
||||
)
|
||||
server.sendResourceListChanged()
|
||||
serverSession.sendResourceListChanged()
|
||||
|
||||
// This should fail because the server doesn't have the tools capability
|
||||
val ex = assertFailsWith<IllegalStateException> {
|
||||
server.sendToolListChanged()
|
||||
serverSession.sendToolListChanged()
|
||||
}
|
||||
assertTrue(ex.message?.contains("Server does not support notifying of tool list changes") == true)
|
||||
}
|
||||
|
|
@ -418,7 +424,29 @@ class ClientTest {
|
|||
|
||||
val def = CompletableDeferred<Unit>()
|
||||
val defTimeOut = CompletableDeferred<Unit>()
|
||||
server.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(capabilities = ClientCapabilities()),
|
||||
)
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
|
||||
serverSession.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
// Simulate delay
|
||||
def.complete(Unit)
|
||||
try {
|
||||
|
|
@ -431,24 +459,6 @@ class ClientTest {
|
|||
fail("Shouldn't have been called")
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(capabilities = ClientCapabilities()),
|
||||
)
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val defCancel = CompletableDeferred<Unit>()
|
||||
val job = launch {
|
||||
try {
|
||||
|
|
@ -478,7 +488,27 @@ class ClientTest {
|
|||
),
|
||||
)
|
||||
|
||||
server.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(capabilities = ClientCapabilities()),
|
||||
)
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
serverSession.setRequestHandler<ListResourcesRequest>(Method.Defined.ResourcesList) { _, _ ->
|
||||
// Simulate a delayed response
|
||||
// Wait ~100ms unless canceled
|
||||
try {
|
||||
|
|
@ -492,23 +522,6 @@ class ClientTest {
|
|||
ListResourcesResult(resources = emptyList())
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
val client = Client(
|
||||
clientInfo = Implementation(name = "test client", version = "1.0"),
|
||||
options = ClientOptions(capabilities = ClientCapabilities()),
|
||||
)
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
// Request with 1 msec timeout should fail immediately
|
||||
val ex = assertFailsWith<Exception> {
|
||||
withTimeout(1) {
|
||||
|
|
@ -559,34 +572,6 @@ class ClientTest {
|
|||
serverOptions,
|
||||
)
|
||||
|
||||
server.setRequestHandler<InitializeRequest>(Method.Defined.Initialize) { _, _ ->
|
||||
InitializeResult(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ServerCapabilities(
|
||||
resources = ServerCapabilities.Resources(null, null),
|
||||
tools = ServerCapabilities.Tools(null),
|
||||
),
|
||||
serverInfo = Implementation(name = "test", version = "1.0"),
|
||||
)
|
||||
}
|
||||
val serverListToolsResult = ListToolsResult(
|
||||
tools = listOf(
|
||||
Tool(
|
||||
name = "testTool",
|
||||
title = "testTool title",
|
||||
description = "testTool description",
|
||||
annotations = null,
|
||||
inputSchema = Tool.Input(),
|
||||
outputSchema = null,
|
||||
),
|
||||
),
|
||||
nextCursor = null,
|
||||
)
|
||||
|
||||
server.setRequestHandler<ListToolsRequest>(Method.Defined.ToolsList) { _, _ ->
|
||||
serverListToolsResult
|
||||
}
|
||||
|
||||
val (clientTransport, serverTransport) = InMemoryTransport.createLinkedPair()
|
||||
|
||||
val client = Client(
|
||||
|
|
@ -601,15 +586,50 @@ class ClientTest {
|
|||
receivedMessage = msg
|
||||
}
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
server.connect(serverTransport)
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
|
||||
serverSession.setRequestHandler<InitializeRequest>(Method.Defined.Initialize) { _, _ ->
|
||||
InitializeResult(
|
||||
protocolVersion = LATEST_PROTOCOL_VERSION,
|
||||
capabilities = ServerCapabilities(
|
||||
resources = ServerCapabilities.Resources(null, null),
|
||||
tools = ServerCapabilities.Tools(null),
|
||||
),
|
||||
serverInfo = Implementation(name = "test", version = "1.0"),
|
||||
)
|
||||
}
|
||||
|
||||
val serverListToolsResult = ListToolsResult(
|
||||
tools = listOf(
|
||||
Tool(
|
||||
name = "testTool",
|
||||
title = "testTool title",
|
||||
description = "testTool description",
|
||||
annotations = null,
|
||||
inputSchema = Tool.Input(),
|
||||
outputSchema = null,
|
||||
),
|
||||
),
|
||||
nextCursor = null,
|
||||
)
|
||||
|
||||
serverSession.setRequestHandler<ListToolsRequest>(Method.Defined.ToolsList) { _, _ ->
|
||||
serverListToolsResult
|
||||
}
|
||||
|
||||
val serverCapabilities = client.serverCapabilities
|
||||
assertEquals(ServerCapabilities.Tools(null), serverCapabilities?.tools)
|
||||
|
||||
|
|
@ -652,15 +672,25 @@ class ClientTest {
|
|||
),
|
||||
)
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch { client.connect(clientTransport) },
|
||||
launch { server.connect(serverTransport) },
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val clientCapabilities = server.clientCapabilities
|
||||
val serverSession = serverSessionResult.await()
|
||||
|
||||
val clientCapabilities = serverSession.clientCapabilities
|
||||
assertEquals(ClientCapabilities.Roots(null), clientCapabilities?.roots)
|
||||
|
||||
val listRootsResult = server.listRoots()
|
||||
val listRootsResult = serverSession.listRoots()
|
||||
|
||||
assertEquals(listRootsResult.roots, clientRoots)
|
||||
}
|
||||
|
|
@ -773,16 +803,28 @@ class ClientTest {
|
|||
|
||||
// Track notifications
|
||||
var rootListChangedNotificationReceived = false
|
||||
server.setNotificationHandler<RootsListChangedNotification>(Method.Defined.NotificationsRootsListChanged) {
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
serverSession.setNotificationHandler<RootsListChangedNotification>(
|
||||
Method.Defined.NotificationsRootsListChanged,
|
||||
) {
|
||||
rootListChangedNotificationReceived = true
|
||||
CompletableDeferred(Unit)
|
||||
}
|
||||
|
||||
listOf(
|
||||
launch { client.connect(clientTransport) },
|
||||
launch { server.connect(serverTransport) },
|
||||
).joinAll()
|
||||
|
||||
client.sendRootsListChanged()
|
||||
|
||||
assertTrue(
|
||||
|
|
@ -809,14 +851,24 @@ class ClientTest {
|
|||
),
|
||||
)
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch { client.connect(clientTransport) },
|
||||
launch { server.connect(serverTransport) },
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val serverSession = serverSessionResult.await()
|
||||
|
||||
// Verify that creating an elicitation throws an exception
|
||||
val exception = assertFailsWith<IllegalStateException> {
|
||||
server.createElicitation(
|
||||
serverSession.createElicitation(
|
||||
message = "Please provide your GitHub username",
|
||||
requestedSchema = CreateElicitationRequest.RequestedSchema(
|
||||
properties = buildJsonObject {
|
||||
|
|
@ -879,12 +931,22 @@ class ClientTest {
|
|||
),
|
||||
)
|
||||
|
||||
val serverSessionResult = CompletableDeferred<ServerSession>()
|
||||
|
||||
listOf(
|
||||
launch { client.connect(clientTransport) },
|
||||
launch { server.connect(serverTransport) },
|
||||
launch {
|
||||
client.connect(clientTransport)
|
||||
println("Client connected")
|
||||
},
|
||||
launch {
|
||||
serverSessionResult.complete(server.connect(serverTransport))
|
||||
println("Server connected")
|
||||
},
|
||||
).joinAll()
|
||||
|
||||
val result = server.createElicitation(
|
||||
val serverSession = serverSessionResult.await()
|
||||
|
||||
val result = serverSession.createElicitation(
|
||||
message = elicitationMessage,
|
||||
requestedSchema = requestedSchema,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -70,19 +70,6 @@ class SseTransportTest : BaseTransportTest() {
|
|||
install(ServerSSE)
|
||||
routing {
|
||||
mcp { mcpServer }
|
||||
// sse {
|
||||
// mcpSseTransport("", transports).apply {
|
||||
// onMessage {
|
||||
// send(it)
|
||||
// }
|
||||
//
|
||||
// start()
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// post {
|
||||
// mcpPostEndpoint(transports)
|
||||
// }
|
||||
}
|
||||
}.startSuspend(wait = false)
|
||||
|
||||
|
|
@ -110,22 +97,7 @@ class SseTransportTest : BaseTransportTest() {
|
|||
val server = embeddedServer(CIO, port = 0) {
|
||||
install(ServerSSE)
|
||||
routing {
|
||||
mcp("/sse") { mcpServer }
|
||||
// route("/sse") {
|
||||
// sse {
|
||||
// mcpSseTransport("", transports).apply {
|
||||
// onMessage {
|
||||
// send(it)
|
||||
// }
|
||||
//
|
||||
// start()
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// post {
|
||||
// mcpPostEndpoint(transports)
|
||||
// }
|
||||
// }
|
||||
mcp { mcpServer }
|
||||
}
|
||||
}.startSuspend(wait = false)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,22 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.integration
|
||||
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.client.plugins.sse.SSE
|
||||
import io.ktor.server.application.install
|
||||
import io.ktor.server.cio.CIOApplicationEngine
|
||||
import io.ktor.server.engine.EmbeddedServer
|
||||
import io.ktor.server.engine.embeddedServer
|
||||
import io.ktor.server.routing.routing
|
||||
import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptArgument
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.Role
|
||||
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.TextContent
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.Client
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.mcpSse
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.mcpSseTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.Server
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.mcp
|
||||
|
|
@ -17,47 +24,172 @@ import kotlinx.coroutines.Dispatchers
|
|||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.coroutines.withContext
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.fail
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
import io.ktor.client.engine.cio.CIO as ClientCIO
|
||||
import io.ktor.client.plugins.sse.SSE as ClientSSE
|
||||
import io.ktor.server.cio.CIO as ServerCIO
|
||||
import io.ktor.server.sse.SSE as ServerSSE
|
||||
|
||||
private const val URL = "127.0.0.1"
|
||||
|
||||
class SseIntegrationTest {
|
||||
@Test
|
||||
fun `client should be able to connect to sse server`() = runTest {
|
||||
val serverEngine = initServer()
|
||||
fun `client should be able to connect to sse server`() = runTest(timeout = 5.seconds) {
|
||||
var server: EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration>? = null
|
||||
var client: Client? = null
|
||||
|
||||
try {
|
||||
withContext(Dispatchers.Default) {
|
||||
val port = serverEngine.engine.resolvedConnectors().first().port
|
||||
val client = initClient(port)
|
||||
client.close()
|
||||
server = initServer()
|
||||
val port = server.engine.resolvedConnectors().first().port
|
||||
client = initClient(serverPort = port)
|
||||
}
|
||||
} catch (e: Exception) {
|
||||
fail("Failed to connect client: $e")
|
||||
} finally {
|
||||
// Make sure to stop the server
|
||||
serverEngine.stopSuspend(1000, 2000)
|
||||
client?.close()
|
||||
server?.stopSuspend(1000, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun initClient(port: Int): Client = HttpClient(ClientCIO) {
|
||||
install(ClientSSE)
|
||||
}.mcpSse("http://$URL:$port")
|
||||
/**
|
||||
* Test Case #1: One opened connection, a client gets a prompt
|
||||
*
|
||||
* 1. Open SSE from Client A.
|
||||
* 2. Send a POST request from Client A to POST /prompts/get.
|
||||
* 3. Observe that Client A receives a response related to it.
|
||||
*/
|
||||
@Test
|
||||
fun `single sse connection`() = runTest(timeout = 5.seconds) {
|
||||
var server: EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration>? = null
|
||||
var client: Client? = null
|
||||
try {
|
||||
withContext(Dispatchers.Default) {
|
||||
server = initServer()
|
||||
val port = server.engine.resolvedConnectors().first().port
|
||||
client = initClient("Client A", port)
|
||||
|
||||
val promptA = getPrompt(client, "Client A")
|
||||
assertTrue { "Client A" in promptA }
|
||||
}
|
||||
} finally {
|
||||
client?.close()
|
||||
server?.stopSuspend(1000, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Test Case #1: Two open connections, each client gets a client-specific prompt
|
||||
*
|
||||
* 1. Open SSE connection #1 from Client A and note the sessionId=<sessionId#1> value.
|
||||
* 2. Open SSE connection #2 from Client B and note the sessionId=<sessionId#2> value.
|
||||
* 3. Send a POST request to POST /message with the corresponding sessionId#1.
|
||||
* 4. Observe that Client B (connection #2) receives a response related to sessionId#1.
|
||||
*/
|
||||
@Test
|
||||
fun `multiple sse connections`() = runTest(timeout = 5.seconds) {
|
||||
var server: EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration>? = null
|
||||
var clientA: Client? = null
|
||||
var clientB: Client? = null
|
||||
|
||||
try {
|
||||
withContext(Dispatchers.Default) {
|
||||
server = initServer()
|
||||
val port = server.engine.resolvedConnectors().first().port
|
||||
|
||||
clientA = initClient("Client A", port)
|
||||
clientB = initClient("Client B", port)
|
||||
|
||||
// Step 3: Send a prompt request from Client A
|
||||
val promptA = getPrompt(clientA, "Client A")
|
||||
// Step 4: Send a prompt request from Client B
|
||||
val promptB = getPrompt(clientB, "Client B")
|
||||
|
||||
assertTrue { "Client A" in promptA }
|
||||
assertTrue { "Client B" in promptB }
|
||||
}
|
||||
} finally {
|
||||
clientA?.close()
|
||||
clientB?.close()
|
||||
server?.stopSuspend(1000, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun initClient(name: String = "", serverPort: Int): Client {
|
||||
val client = Client(
|
||||
Implementation(name = name, version = "1.0.0"),
|
||||
)
|
||||
|
||||
val httpClient = HttpClient(ClientCIO) {
|
||||
install(SSE)
|
||||
}
|
||||
|
||||
// Create a transport wrapper that captures the session ID and received messages
|
||||
val transport = httpClient.mcpSseTransport {
|
||||
url {
|
||||
host = URL
|
||||
port = serverPort
|
||||
}
|
||||
}
|
||||
|
||||
client.connect(transport)
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
private suspend fun initServer(): EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration> {
|
||||
val server = Server(
|
||||
Implementation(name = "sse-e2e-test", version = "1.0.0"),
|
||||
ServerOptions(capabilities = ServerCapabilities()),
|
||||
Implementation(name = "sse-server", version = "1.0.0"),
|
||||
ServerOptions(
|
||||
capabilities = ServerCapabilities(prompts = ServerCapabilities.Prompts(listChanged = true)),
|
||||
),
|
||||
)
|
||||
|
||||
return embeddedServer(ServerCIO, host = URL, port = 0) {
|
||||
server.addPrompt(
|
||||
name = "prompt",
|
||||
description = "Prompt description",
|
||||
arguments = listOf(
|
||||
PromptArgument(
|
||||
name = "client",
|
||||
description = "Client name who requested a prompt",
|
||||
required = true,
|
||||
),
|
||||
),
|
||||
) { request ->
|
||||
GetPromptResult(
|
||||
"Prompt for ${request.name}",
|
||||
messages = listOf(
|
||||
PromptMessage(
|
||||
role = Role.user,
|
||||
content = TextContent("Prompt for client ${request.arguments?.get("client")}"),
|
||||
),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
val ktorServer = embeddedServer(ServerCIO, host = URL, port = PORT) {
|
||||
install(ServerSSE)
|
||||
routing {
|
||||
mcp { server }
|
||||
}
|
||||
}.startSuspend(wait = false)
|
||||
}
|
||||
|
||||
return ktorServer.startSuspend(wait = false)
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves a prompt result using the provided client and client name.
|
||||
*/
|
||||
private suspend fun getPrompt(client: Client, clientName: String): String {
|
||||
val response = client.getPrompt(
|
||||
GetPromptRequest(
|
||||
"prompt",
|
||||
arguments = mapOf("client" to clientName),
|
||||
),
|
||||
)
|
||||
|
||||
return (response.messages.first().content as? TextContent)?.text
|
||||
?: error("Failed to receive prompt for Client $clientName")
|
||||
}
|
||||
|
||||
companion object {
|
||||
private const val URL = "127.0.0.1"
|
||||
private const val PORT = 0
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,196 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.integration
|
||||
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.server.application.install
|
||||
import io.ktor.server.cio.CIOApplicationEngine
|
||||
import io.ktor.server.engine.EmbeddedServer
|
||||
import io.ktor.server.engine.embeddedServer
|
||||
import io.ktor.server.routing.routing
|
||||
import io.modelcontextprotocol.kotlin.sdk.GetPromptRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.GetPromptResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.Implementation
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptArgument
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.Role
|
||||
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
|
||||
import io.modelcontextprotocol.kotlin.sdk.TextContent
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.Client
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.mcpWebSocketTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.Server
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
|
||||
import io.modelcontextprotocol.kotlin.sdk.server.mcpWebSocket
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.coroutines.withContext
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
import io.ktor.client.engine.cio.CIO as ClientCIO
|
||||
import io.ktor.client.plugins.websocket.WebSockets as ClientWebSocket
|
||||
import io.ktor.server.cio.CIO as ServerCIO
|
||||
import io.ktor.server.websocket.WebSockets as ServerWebSockets
|
||||
|
||||
class WebSocketIntegrationTest {
|
||||
|
||||
@Test
|
||||
fun `client should be able to connect to websocket server 2`() = runTest(timeout = 5.seconds) {
|
||||
var server: EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration>? = null
|
||||
var client: Client? = null
|
||||
|
||||
try {
|
||||
withContext(Dispatchers.Default) {
|
||||
server = initServer()
|
||||
val port = server.engine.resolvedConnectors().first().port
|
||||
client = initClient(serverPort = port)
|
||||
}
|
||||
} finally {
|
||||
client?.close()
|
||||
server?.stopSuspend(1000, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Test Case #1: One opened connection, a client gets a prompt
|
||||
*
|
||||
* 1. Open WebSocket from Client A.
|
||||
* 2. Send a POST request from Client A to POST /prompts/get.
|
||||
* 3. Observe that Client A receives a response related to it.
|
||||
*/
|
||||
@Test
|
||||
fun `single websocket connection`() = runTest(timeout = 5.seconds) {
|
||||
var server: EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration>? = null
|
||||
var client: Client? = null
|
||||
|
||||
try {
|
||||
withContext(Dispatchers.Default) {
|
||||
server = initServer()
|
||||
val port = server.engine.resolvedConnectors().first().port
|
||||
client = initClient("Client A", port)
|
||||
|
||||
val promptA = getPrompt(client, "Client A")
|
||||
assertTrue { "Client A" in promptA }
|
||||
}
|
||||
} finally {
|
||||
client?.close()
|
||||
server?.stopSuspend(1000, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Test Case #1: Two open connections, each client gets a client-specific prompt
|
||||
*
|
||||
* 1. Open WebSocket connection #1 from Client A and note the sessionId=<sessionId#1> value.
|
||||
* 2. Open WebSocket connection #2 from Client B and note the sessionId=<sessionId#2> value.
|
||||
* 3. Send a POST request to POST /message with the corresponding sessionId#1.
|
||||
* 4. Observe that Client B (connection #2) receives a response related to sessionId#1.
|
||||
*/
|
||||
@Test
|
||||
fun `multiple websocket connections`() = runTest(timeout = 5.seconds) {
|
||||
var server: EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration>? = null
|
||||
var clientA: Client? = null
|
||||
var clientB: Client? = null
|
||||
|
||||
try {
|
||||
withContext(Dispatchers.Default) {
|
||||
server = initServer()
|
||||
val port = server.engine.resolvedConnectors().first().port
|
||||
clientA = initClient("Client A", port)
|
||||
clientB = initClient("Client B", port)
|
||||
|
||||
// Step 3: Send a prompt request from Client A
|
||||
val promptA = getPrompt(clientA, "Client A")
|
||||
// Step 4: Send a prompt request from Client B
|
||||
val promptB = getPrompt(clientB, "Client B")
|
||||
|
||||
assertTrue { "Client A" in promptA }
|
||||
assertTrue { "Client B" in promptB }
|
||||
}
|
||||
} finally {
|
||||
clientA?.close()
|
||||
clientB?.close()
|
||||
server?.stopSuspend(1000, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun initClient(name: String = "", serverPort: Int): Client {
|
||||
val client = Client(
|
||||
Implementation(name = name, version = "1.0.0"),
|
||||
)
|
||||
|
||||
val httpClient = HttpClient(ClientCIO) {
|
||||
install(ClientWebSocket)
|
||||
}
|
||||
|
||||
// Create a transport wrapper that captures the session ID and received messages
|
||||
val transport = httpClient.mcpWebSocketTransport {
|
||||
url {
|
||||
host = URL
|
||||
port = serverPort
|
||||
}
|
||||
}
|
||||
|
||||
client.connect(transport)
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
private suspend fun initServer(): EmbeddedServer<CIOApplicationEngine, CIOApplicationEngine.Configuration> {
|
||||
val server = Server(
|
||||
Implementation(name = "websocket-server", version = "1.0.0"),
|
||||
ServerOptions(
|
||||
capabilities = ServerCapabilities(prompts = ServerCapabilities.Prompts(listChanged = true)),
|
||||
),
|
||||
)
|
||||
|
||||
server.addPrompt(
|
||||
name = "prompt",
|
||||
description = "Prompt description",
|
||||
arguments = listOf(
|
||||
PromptArgument(
|
||||
name = "client",
|
||||
description = "Client name who requested a prompt",
|
||||
required = true,
|
||||
),
|
||||
),
|
||||
) { request ->
|
||||
GetPromptResult(
|
||||
"Prompt for ${request.name}",
|
||||
messages = listOf(
|
||||
PromptMessage(
|
||||
role = Role.user,
|
||||
content = TextContent("Prompt for client ${request.arguments?.get("client")}"),
|
||||
),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
val ktorServer = embeddedServer(ServerCIO, host = URL, port = PORT) {
|
||||
install(ServerWebSockets)
|
||||
routing {
|
||||
mcpWebSocket(block = { server })
|
||||
}
|
||||
}
|
||||
|
||||
return ktorServer.startSuspend(wait = false)
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves a prompt result using the provided client and client name.
|
||||
*/
|
||||
private suspend fun getPrompt(client: Client, clientName: String): String {
|
||||
val response = client.getPrompt(
|
||||
GetPromptRequest(
|
||||
"prompt",
|
||||
arguments = mapOf("client" to clientName),
|
||||
),
|
||||
)
|
||||
|
||||
return (response.messages.first().content as? TextContent)?.text
|
||||
?: error("Failed to receive prompt for Client $clientName")
|
||||
}
|
||||
|
||||
companion object {
|
||||
private const val PORT = 0
|
||||
private const val URL = "127.0.0.1"
|
||||
}
|
||||
}
|
||||
|
|
@ -1,8 +1,6 @@
|
|||
package io.modelcontextprotocol.kotlin.sdk.integration.kotlin
|
||||
|
||||
import io.modelcontextprotocol.kotlin.sdk.BlobResourceContents
|
||||
import io.modelcontextprotocol.kotlin.sdk.EmptyRequestResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.Method
|
||||
import io.modelcontextprotocol.kotlin.sdk.ReadResourceRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
|
||||
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
|
||||
|
|
@ -16,6 +14,7 @@ import kotlinx.coroutines.test.runTest
|
|||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.api.assertThrows
|
||||
import java.util.concurrent.atomic.AtomicBoolean
|
||||
import kotlin.test.Ignore
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertTrue
|
||||
|
|
@ -67,14 +66,6 @@ abstract class AbstractResourceIntegrationTest : KotlinTestBase() {
|
|||
)
|
||||
}
|
||||
|
||||
server.setRequestHandler<SubscribeRequest>(Method.Defined.ResourcesSubscribe) { _, _ ->
|
||||
EmptyRequestResult()
|
||||
}
|
||||
|
||||
server.setRequestHandler<UnsubscribeRequest>(Method.Defined.ResourcesUnsubscribe) { _, _ ->
|
||||
EmptyRequestResult()
|
||||
}
|
||||
|
||||
server.addResource(
|
||||
uri = testResourceUri,
|
||||
name = testResourceName,
|
||||
|
|
@ -169,6 +160,7 @@ abstract class AbstractResourceIntegrationTest : KotlinTestBase() {
|
|||
assertEquals(testResourceContent, content.text, "Resource content should match")
|
||||
}
|
||||
|
||||
@Ignore("Blocked by https://github.com/modelcontextprotocol/kotlin-sdk/issues/249")
|
||||
@Test
|
||||
fun testSubscribeAndUnsubscribe() {
|
||||
runBlocking(Dispatchers.IO) {
|
||||
|
|
|
|||
|
|
@ -26,8 +26,6 @@ import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
|||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCResponse
|
||||
import io.modelcontextprotocol.kotlin.sdk.LoggingLevel
|
||||
import io.modelcontextprotocol.kotlin.sdk.LoggingMessageNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptArgument
|
||||
import io.modelcontextprotocol.kotlin.sdk.PromptMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.ReadResourceResult
|
||||
|
|
@ -254,32 +252,6 @@ class KotlinServerForTsClient {
|
|||
) { request ->
|
||||
val name = (request.arguments["name"] as? JsonPrimitive)?.content ?: "World"
|
||||
|
||||
server.sendToolListChanged()
|
||||
server.sendLoggingMessage(
|
||||
LoggingMessageNotification(
|
||||
LoggingMessageNotification.Params(
|
||||
level = LoggingLevel.info,
|
||||
data = JsonPrimitive("Preparing greeting for $name"),
|
||||
),
|
||||
),
|
||||
)
|
||||
server.sendLoggingMessage(
|
||||
LoggingMessageNotification(
|
||||
LoggingMessageNotification.Params(
|
||||
level = LoggingLevel.info,
|
||||
data = JsonPrimitive("Halfway there for $name"),
|
||||
),
|
||||
),
|
||||
)
|
||||
server.sendLoggingMessage(
|
||||
LoggingMessageNotification(
|
||||
LoggingMessageNotification.Params(
|
||||
level = LoggingLevel.info,
|
||||
data = JsonPrimitive("Done sending greetings to $name"),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
CallToolResult(
|
||||
content = listOf(TextContent("Multiple greetings sent to $name!")),
|
||||
structuredContent = buildJsonObject {
|
||||
|
|
|
|||
Loading…
Reference in a new issue