Add streamable http client transport (#147)
* feat: add streamable http client * update StreamableHttpClientTransport * improve StreamableHttpClientTransport handling of headers, event parsing, and session termination * add simple tests for StreamableClient and fix send * add client mocking * add notification e2e tests --------- Co-authored-by: Zach Tang <zachtangoj@gmail.com>
This commit is contained in:
parent
12f686fa4b
commit
2a38e59bd3
1 changed files with 788 additions and 1 deletions
369
src/jvmTest/kotlin/client/StreamableHttpClientTransportTest.kt
Normal file
369
src/jvmTest/kotlin/client/StreamableHttpClientTransportTest.kt
Normal file
|
|
@ -0,0 +1,369 @@
|
|||
package client
|
||||
|
||||
import io.ktor.client.HttpClient
|
||||
import io.ktor.client.engine.mock.MockEngine
|
||||
import io.ktor.client.engine.mock.MockRequestHandler
|
||||
import io.ktor.client.engine.mock.respond
|
||||
import io.ktor.client.plugins.sse.SSE
|
||||
import io.ktor.http.ContentType
|
||||
import io.ktor.http.HttpHeaders
|
||||
import io.ktor.http.HttpMethod
|
||||
import io.ktor.http.HttpStatusCode
|
||||
import io.ktor.http.content.TextContent
|
||||
import io.ktor.http.headersOf
|
||||
import io.ktor.utils.io.ByteReadChannel
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCNotification
|
||||
import io.modelcontextprotocol.kotlin.sdk.JSONRPCRequest
|
||||
import io.modelcontextprotocol.kotlin.sdk.RequestId
|
||||
import io.modelcontextprotocol.kotlin.sdk.client.StreamableHttpClientTransport
|
||||
import io.modelcontextprotocol.kotlin.sdk.shared.McpJson
|
||||
import kotlinx.coroutines.delay
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonPrimitive
|
||||
import kotlinx.serialization.json.buildJsonObject
|
||||
import org.junit.jupiter.api.assertDoesNotThrow
|
||||
import kotlin.test.Ignore
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNull
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.time.Duration.Companion.seconds
|
||||
|
||||
class StreamableHttpClientTransportTest {
|
||||
|
||||
private fun createTransport(handler: MockRequestHandler): StreamableHttpClientTransport {
|
||||
val mockEngine = MockEngine(handler)
|
||||
val httpClient = HttpClient(mockEngine) {
|
||||
install(SSE) {
|
||||
reconnectionTime = 1.seconds
|
||||
}
|
||||
}
|
||||
|
||||
return StreamableHttpClientTransport(httpClient, url = "http://localhost:8080/mcp")
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testSendJsonRpcMessage() = runTest {
|
||||
val message = JSONRPCRequest(
|
||||
id = RequestId.StringId("test-id"),
|
||||
method = "test",
|
||||
params = buildJsonObject { }
|
||||
)
|
||||
|
||||
val transport = createTransport { request ->
|
||||
assertEquals(HttpMethod.Post, request.method)
|
||||
assertEquals("http://localhost:8080/mcp", request.url.toString())
|
||||
assertEquals(ContentType.Application.Json, request.body.contentType)
|
||||
|
||||
val body = (request.body as TextContent).text
|
||||
val decodedMessage = McpJson.decodeFromString<JSONRPCMessage>(body)
|
||||
assertEquals(message, decodedMessage)
|
||||
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.Accepted
|
||||
)
|
||||
}
|
||||
|
||||
transport.start()
|
||||
transport.send(message)
|
||||
transport.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testStoreSessionId() = runTest {
|
||||
val initMessage = JSONRPCRequest(
|
||||
id = RequestId.StringId("test-id"),
|
||||
method = "initialize",
|
||||
params = buildJsonObject {
|
||||
put("clientInfo", buildJsonObject {
|
||||
put("name", JsonPrimitive("test-client"))
|
||||
put("version", JsonPrimitive("1.0"))
|
||||
})
|
||||
put("protocolVersion", JsonPrimitive("2025-06-18"))
|
||||
}
|
||||
)
|
||||
|
||||
val transport = createTransport { request ->
|
||||
when (val msg = McpJson.decodeFromString<JSONRPCMessage>((request.body as TextContent).text)) {
|
||||
is JSONRPCRequest if msg.method == "initialize" -> respond(
|
||||
content = "", status = HttpStatusCode.OK,
|
||||
headers = headersOf("mcp-session-id", "test-session-id")
|
||||
)
|
||||
|
||||
is JSONRPCNotification if msg.method == "test" -> {
|
||||
assertEquals("test-session-id", request.headers["mcp-session-id"])
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.Accepted
|
||||
)
|
||||
}
|
||||
|
||||
else -> error("Unexpected message: $msg")
|
||||
}
|
||||
}
|
||||
|
||||
transport.start()
|
||||
transport.send(initMessage)
|
||||
|
||||
assertEquals("test-session-id", transport.sessionId)
|
||||
|
||||
transport.send(JSONRPCNotification(method = "test"))
|
||||
|
||||
transport.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testTerminateSession() = runTest {
|
||||
// transport.sessionId = "test-session-id"
|
||||
|
||||
val transport = createTransport { request ->
|
||||
assertEquals(HttpMethod.Delete, request.method)
|
||||
assertEquals("test-session-id", request.headers["mcp-session-id"])
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.OK
|
||||
)
|
||||
}
|
||||
|
||||
transport.start()
|
||||
transport.terminateSession()
|
||||
|
||||
assertNull(transport.sessionId)
|
||||
transport.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testTerminateSessionHandle405() = runTest {
|
||||
// transport.sessionId = "test-session-id"
|
||||
|
||||
val transport = createTransport { request ->
|
||||
assertEquals(HttpMethod.Delete, request.method)
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.MethodNotAllowed
|
||||
)
|
||||
}
|
||||
|
||||
transport.start()
|
||||
// Should not throw for 405
|
||||
assertDoesNotThrow {
|
||||
transport.terminateSession()
|
||||
}
|
||||
|
||||
// Session ID should still be cleared
|
||||
assertNull(transport.sessionId)
|
||||
transport.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testProtocolVersionHeader() = runTest {
|
||||
val transport = createTransport { request ->
|
||||
assertEquals("2025-06-18", request.headers["mcp-protocol-version"])
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.Accepted
|
||||
)
|
||||
}
|
||||
transport.protocolVersion = "2025-06-18"
|
||||
|
||||
transport.start()
|
||||
transport.send(JSONRPCNotification(method = "test"))
|
||||
transport.close()
|
||||
}
|
||||
|
||||
@Ignore("Engine doesn't support SSECapability: https://youtrack.jetbrains.com/issue/KTOR-8177/MockEngine-Add-SSE-support")
|
||||
@Test
|
||||
fun testNotificationSchemaE2E() = runTest {
|
||||
val receivedMessages = mutableListOf<JSONRPCMessage>()
|
||||
var sseStarted = false
|
||||
|
||||
val transport = createTransport { request ->
|
||||
when (request.method) {
|
||||
HttpMethod.Post if request.body.toString().contains("notifications/initialized") -> {
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.Accepted,
|
||||
headers = headersOf("mcp-session-id", "notification-test-session")
|
||||
)
|
||||
}
|
||||
|
||||
// Handle SSE connection
|
||||
HttpMethod.Get -> {
|
||||
sseStarted = true
|
||||
val sseContent = buildString {
|
||||
// Server sends various notifications
|
||||
appendLine("event: message")
|
||||
appendLine("id: 1")
|
||||
appendLine("""data: {"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":"upload-123","progress":50,"total":100}}""")
|
||||
appendLine()
|
||||
|
||||
appendLine("event: message")
|
||||
appendLine("id: 2")
|
||||
appendLine("""data: {"jsonrpc":"2.0","method":"notifications/resources/list_changed"}""")
|
||||
appendLine()
|
||||
|
||||
appendLine("event: message")
|
||||
appendLine("id: 3")
|
||||
appendLine("""data: {"jsonrpc":"2.0","method":"notifications/tools/list_changed"}""")
|
||||
appendLine()
|
||||
}
|
||||
respond(
|
||||
content = ByteReadChannel(sseContent),
|
||||
status = HttpStatusCode.OK,
|
||||
headers = headersOf(
|
||||
HttpHeaders.ContentType, ContentType.Text.EventStream.toString()
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
// Handle regular notifications
|
||||
HttpMethod.Post -> {
|
||||
respond(
|
||||
content = "",
|
||||
status = HttpStatusCode.Accepted
|
||||
)
|
||||
}
|
||||
|
||||
else -> respond("", HttpStatusCode.OK)
|
||||
}
|
||||
}
|
||||
|
||||
transport.onMessage { message ->
|
||||
receivedMessages.add(message)
|
||||
}
|
||||
|
||||
transport.start()
|
||||
|
||||
// Test 1: Send initialized notification to trigger SSE
|
||||
val initializedNotification = JSONRPCNotification(
|
||||
method = "notifications/initialized",
|
||||
params = buildJsonObject {
|
||||
put("protocolVersion", JsonPrimitive("1.0"))
|
||||
put("capabilities", buildJsonObject {
|
||||
put("tools", JsonPrimitive(true))
|
||||
put("resources", JsonPrimitive(true))
|
||||
})
|
||||
}
|
||||
)
|
||||
|
||||
transport.send(initializedNotification)
|
||||
|
||||
// Verify SSE was triggered
|
||||
assertTrue(sseStarted, "SSE should start after initialized notification")
|
||||
|
||||
// Test 2: Verify received notifications
|
||||
assertEquals(3, receivedMessages.size)
|
||||
assertTrue(receivedMessages.all { it is JSONRPCNotification })
|
||||
|
||||
val notifications = receivedMessages.filterIsInstance<JSONRPCNotification>()
|
||||
|
||||
// Verify progress notification
|
||||
val progressNotif = notifications[0]
|
||||
assertEquals("notifications/progress", progressNotif.method)
|
||||
val progressParams = progressNotif.params as JsonObject
|
||||
assertEquals("upload-123", (progressParams["progressToken"] as JsonPrimitive).content)
|
||||
assertEquals(50, (progressParams["progress"] as JsonPrimitive).content.toInt())
|
||||
|
||||
// Verify list changed notifications
|
||||
assertEquals("notifications/resources/list_changed", notifications[1].method)
|
||||
assertEquals("notifications/tools/list_changed", notifications[2].method)
|
||||
|
||||
// Test 3: Send various client notifications
|
||||
val clientNotifications = listOf(
|
||||
JSONRPCNotification(
|
||||
method = "notifications/progress",
|
||||
params = buildJsonObject {
|
||||
put("progressToken", JsonPrimitive("download-456"))
|
||||
put("progress", JsonPrimitive(75))
|
||||
}
|
||||
),
|
||||
JSONRPCNotification(
|
||||
method = "notifications/cancelled",
|
||||
params = buildJsonObject {
|
||||
put("requestId", JsonPrimitive("req-789"))
|
||||
put("reason", JsonPrimitive("user_cancelled"))
|
||||
}
|
||||
),
|
||||
JSONRPCNotification(
|
||||
method = "notifications/message",
|
||||
params = buildJsonObject {
|
||||
put("level", JsonPrimitive("info"))
|
||||
put("message", JsonPrimitive("Operation completed"))
|
||||
put("data", buildJsonObject {
|
||||
put("duration", JsonPrimitive(1234))
|
||||
})
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
// Send all client notifications
|
||||
clientNotifications.forEach { notification ->
|
||||
transport.send(notification)
|
||||
}
|
||||
|
||||
// Verify session ID is maintained
|
||||
assertEquals("notification-test-session", transport.sessionId)
|
||||
transport.close()
|
||||
}
|
||||
|
||||
@Ignore("Engine doesn't support SSECapability: https://youtrack.jetbrains.com/issue/KTOR-8177/MockEngine-Add-SSE-support")
|
||||
@Test
|
||||
fun testNotificationWithResumptionToken() = runTest {
|
||||
var resumptionTokenReceived: String? = null
|
||||
var lastEventIdSent: String? = null
|
||||
|
||||
val transport = createTransport { request ->
|
||||
// Capture Last-Event-ID header
|
||||
lastEventIdSent = request.headers["Last-Event-ID"]
|
||||
|
||||
when (request.method) {
|
||||
HttpMethod.Get -> {
|
||||
val sseContent = buildString {
|
||||
appendLine("event: message")
|
||||
appendLine("id: resume-100")
|
||||
appendLine("""data: {"jsonrpc":"2.0","method":"notifications/resumed","params":{"fromToken":"${lastEventIdSent}"}}""")
|
||||
appendLine()
|
||||
}
|
||||
respond(
|
||||
content = ByteReadChannel(sseContent),
|
||||
status = HttpStatusCode.OK,
|
||||
headers = headersOf(
|
||||
HttpHeaders.ContentType, ContentType.Text.EventStream.toString()
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
else -> respond("", HttpStatusCode.Accepted)
|
||||
}
|
||||
}
|
||||
|
||||
transport.start()
|
||||
|
||||
// Send notification with resumption token
|
||||
transport.send(
|
||||
message = JSONRPCNotification(
|
||||
method = "notifications/test",
|
||||
params = buildJsonObject {
|
||||
put("data", JsonPrimitive("test-data"))
|
||||
}
|
||||
),
|
||||
resumptionToken = "previous-token-99",
|
||||
onResumptionToken = { token ->
|
||||
resumptionTokenReceived = token
|
||||
}
|
||||
)
|
||||
|
||||
// Wait for response
|
||||
delay(1.seconds)
|
||||
|
||||
// Verify resumption token was sent in header
|
||||
assertEquals("previous-token-99", lastEventIdSent)
|
||||
|
||||
// Verify new resumption token was received
|
||||
assertEquals("resume-100", resumptionTokenReceived)
|
||||
transport.close()
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue