fec680a243* Prevent double continuation resumption in client Consolidate request removal and continuation resumption logic to ensure each request's continuation is resumed exactly once, preventing "SWIFT TASK CONTINUATION MISUSE" errors during network failures. * Fix race condition test
564 lines
20 KiB
Swift
564 lines
20 KiB
Swift
import Foundation
|
|
import Testing
|
|
|
|
@testable import MCP
|
|
|
|
@Suite("Client Tests")
|
|
struct ClientTests {
|
|
@Test("Client connect and disconnect")
|
|
func testClientConnectAndDisconnect() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
#expect(await transport.isConnected == false)
|
|
try await client.connect(transport: transport)
|
|
#expect(await transport.isConnected == true)
|
|
await client.disconnect()
|
|
#expect(await transport.isConnected == false)
|
|
}
|
|
|
|
@Test(
|
|
"Initialize request",
|
|
.timeLimit(.minutes(1))
|
|
)
|
|
func testClientInitialize() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
// Small delay to ensure message loop is started
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Create a task for initialize that we'll cancel
|
|
let initTask = Task {
|
|
try await client.initialize()
|
|
}
|
|
|
|
// Give it a moment to send the request
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
#expect(await transport.sentMessages.count == 1)
|
|
#expect(await transport.sentMessages.first?.contains(Initialize.name) == true)
|
|
#expect(await transport.sentMessages.first?.contains(client.name) == true)
|
|
#expect(await transport.sentMessages.first?.contains(client.version) == true)
|
|
|
|
// Cancel the initialize task
|
|
initTask.cancel()
|
|
|
|
// Disconnect client to clean up message loop and give time for continuation cleanup
|
|
await client.disconnect()
|
|
try await Task.sleep(for: .milliseconds(50))
|
|
}
|
|
|
|
@Test(
|
|
"Ping request",
|
|
.timeLimit(.minutes(1))
|
|
)
|
|
func testClientPing() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
// Small delay to ensure message loop is started
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Create a task for the ping that we'll cancel
|
|
let pingTask = Task {
|
|
try await client.ping()
|
|
}
|
|
|
|
// Give it a moment to send the request
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
#expect(await transport.sentMessages.count == 1)
|
|
#expect(await transport.sentMessages.first?.contains(Ping.name) == true)
|
|
|
|
// Cancel the ping task
|
|
pingTask.cancel()
|
|
|
|
// Disconnect client to clean up message loop and give time for continuation cleanup
|
|
await client.disconnect()
|
|
try await Task.sleep(for: .milliseconds(50))
|
|
}
|
|
|
|
@Test("Connection failure handling")
|
|
func testClientConnectionFailure() async {
|
|
let transport = MockTransport()
|
|
await transport.setFailConnect(true)
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
do {
|
|
try await client.connect(transport: transport)
|
|
#expect(Bool(false), "Expected connection to fail")
|
|
} catch let error as MCPError {
|
|
if case MCPError.transportError = error {
|
|
#expect(Bool(true))
|
|
} else {
|
|
#expect(Bool(false), "Expected transport error")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCP.Error")
|
|
}
|
|
}
|
|
|
|
@Test("Send failure handling")
|
|
func testClientSendFailure() async throws {
|
|
let transport = MockTransport()
|
|
await transport.setFailSend(true)
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
|
|
do {
|
|
try await client.ping()
|
|
#expect(Bool(false), "Expected ping to fail")
|
|
} catch let error as MCPError {
|
|
if case MCPError.transportError = error {
|
|
#expect(Bool(true))
|
|
} else {
|
|
#expect(Bool(false), "Expected transport error")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCP.Error")
|
|
}
|
|
}
|
|
|
|
@Test("Strict configuration - capabilities check")
|
|
func testStrictConfiguration() async throws {
|
|
let transport = MockTransport()
|
|
let config = Client.Configuration.strict
|
|
let client = Client(name: "TestClient", version: "1.0", configuration: config)
|
|
|
|
try await client.connect(transport: transport)
|
|
|
|
// Create a task for listPrompts
|
|
let promptsTask = Task<Void, Swift.Error> {
|
|
do {
|
|
_ = try await client.listPrompts()
|
|
#expect(Bool(false), "Expected listPrompts to fail in strict mode")
|
|
} catch let error as MCPError {
|
|
if case MCPError.methodNotFound = error {
|
|
#expect(Bool(true))
|
|
} else {
|
|
#expect(Bool(false), "Expected methodNotFound error, got \(error)")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCP.Error")
|
|
}
|
|
}
|
|
|
|
// Give it a short time to execute the task
|
|
try await Task.sleep(for: .milliseconds(50))
|
|
|
|
// Cancel the task if it's still running
|
|
promptsTask.cancel()
|
|
|
|
// Disconnect client
|
|
await client.disconnect()
|
|
try await Task.sleep(for: .milliseconds(50))
|
|
}
|
|
|
|
@Test("Non-strict configuration - capabilities check")
|
|
func testNonStrictConfiguration() async throws {
|
|
let transport = MockTransport()
|
|
let config = Client.Configuration.default
|
|
let client = Client(name: "TestClient", version: "1.0", configuration: config)
|
|
|
|
try await client.connect(transport: transport)
|
|
|
|
// Wait a bit for any setup to complete
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Send the listPrompts request and immediately provide an error response
|
|
let promptsTask = Task {
|
|
do {
|
|
// Start the request
|
|
try await Task.sleep(for: .seconds(1))
|
|
|
|
// Get the last sent message and extract the request ID
|
|
if let lastMessage = await transport.sentMessages.last,
|
|
let data = lastMessage.data(using: .utf8),
|
|
let decodedRequest = try? JSONDecoder().decode(
|
|
Request<ListPrompts>.self, from: data)
|
|
{
|
|
|
|
// Create an error response with the same ID
|
|
let errorResponse = Response<ListPrompts>(
|
|
id: decodedRequest.id,
|
|
error: MCPError.methodNotFound("Test: Prompts capability not available")
|
|
)
|
|
try await transport.queue(response: errorResponse)
|
|
|
|
// Try the request now that we have a response queued
|
|
do {
|
|
_ = try await client.listPrompts()
|
|
#expect(Bool(false), "Expected listPrompts to fail in non-strict mode")
|
|
} catch let error as MCPError {
|
|
if case MCPError.methodNotFound = error {
|
|
#expect(Bool(true))
|
|
} else {
|
|
#expect(Bool(false), "Expected methodNotFound error, got \(error)")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCP.Error")
|
|
}
|
|
}
|
|
} catch {
|
|
// Ignore task cancellation
|
|
if !(error is CancellationError) {
|
|
throw error
|
|
}
|
|
}
|
|
}
|
|
|
|
// Wait for the task to complete or timeout
|
|
let timeoutTask = Task {
|
|
try await Task.sleep(for: .milliseconds(500))
|
|
promptsTask.cancel()
|
|
}
|
|
|
|
// Wait for the task to complete
|
|
_ = await promptsTask.result
|
|
|
|
// Cancel the timeout task
|
|
timeoutTask.cancel()
|
|
|
|
// Disconnect client
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Batch request - success")
|
|
func testBatchRequestSuccess() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10)) // Allow connection tasks
|
|
|
|
let request1 = Ping.request()
|
|
let request2 = Ping.request()
|
|
var resultTask1: Task<Ping.Result, Swift.Error>?
|
|
var resultTask2: Task<Ping.Result, Swift.Error>?
|
|
|
|
try await client.withBatch { batch in
|
|
resultTask1 = try await batch.addRequest(request1)
|
|
resultTask2 = try await batch.addRequest(request2)
|
|
}
|
|
|
|
// Check if one batch message was sent
|
|
let sentMessages = await transport.sentMessages
|
|
#expect(sentMessages.count == 1)
|
|
|
|
guard let batchData = sentMessages.first?.data(using: .utf8) else {
|
|
#expect(Bool(false), "Failed to get batch data")
|
|
return
|
|
}
|
|
|
|
// Verify the sent batch contains the two requests
|
|
let decoder = JSONDecoder()
|
|
let sentRequests = try decoder.decode([AnyRequest].self, from: batchData)
|
|
#expect(sentRequests.count == 2)
|
|
#expect(sentRequests.first?.id == request1.id)
|
|
#expect(sentRequests.first?.method == Ping.name)
|
|
#expect(sentRequests.last?.id == request2.id)
|
|
#expect(sentRequests.last?.method == Ping.name)
|
|
|
|
// Prepare batch response
|
|
let response1 = Response<Ping>(id: request1.id, result: .init())
|
|
let response2 = Response<Ping>(id: request2.id, result: .init())
|
|
let anyResponse1 = try AnyResponse(response1)
|
|
let anyResponse2 = try AnyResponse(response2)
|
|
|
|
// Queue the batch response
|
|
try await transport.queue(batch: [anyResponse1, anyResponse2])
|
|
|
|
// Wait for results and verify
|
|
guard let task1 = resultTask1, let task2 = resultTask2 else {
|
|
#expect(Bool(false), "Result tasks not created")
|
|
return
|
|
}
|
|
|
|
_ = try await task1.value // Should succeed
|
|
_ = try await task2.value // Should succeed
|
|
|
|
#expect(Bool(true)) // Reaching here means success
|
|
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Batch request - mixed success/error")
|
|
func testBatchRequestMixed() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
let request1 = Ping.request() // Success
|
|
let request2 = Ping.request() // Error
|
|
|
|
var resultTasks: [Task<Ping.Result, Swift.Error>] = []
|
|
|
|
try await client.withBatch { batch in
|
|
resultTasks.append(try await batch.addRequest(request1))
|
|
resultTasks.append(try await batch.addRequest(request2))
|
|
}
|
|
|
|
// Check if one batch message was sent
|
|
#expect(await transport.sentMessages.count == 1)
|
|
|
|
// Prepare batch response (success for 1, error for 2)
|
|
let response1 = Response<Ping>(id: request1.id, result: .init())
|
|
let error = MCPError.internalError("Simulated batch error")
|
|
let response2 = Response<Ping>(id: request2.id, error: error)
|
|
let anyResponse1 = try AnyResponse(response1)
|
|
let anyResponse2 = try AnyResponse(response2)
|
|
|
|
// Queue the batch response
|
|
try await transport.queue(batch: [anyResponse1, anyResponse2])
|
|
|
|
// Wait for results and verify
|
|
#expect(resultTasks.count == 2)
|
|
guard resultTasks.count == 2 else {
|
|
#expect(Bool(false), "Expected 2 result tasks")
|
|
return
|
|
}
|
|
|
|
let task1 = resultTasks[0]
|
|
let task2 = resultTasks[1]
|
|
|
|
_ = try await task1.value // Task 1 should succeed
|
|
|
|
do {
|
|
_ = try await task2.value // Task 2 should fail
|
|
#expect(Bool(false), "Task 2 should have thrown an error")
|
|
} catch let mcpError as MCPError {
|
|
if case .internalError(let message) = mcpError {
|
|
#expect(message == "Simulated batch error")
|
|
} else {
|
|
#expect(Bool(false), "Expected internalError, got \(mcpError)")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCPError, got \(error)")
|
|
}
|
|
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Batch request - empty")
|
|
func testBatchRequestEmpty() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Call withBatch but don't add any requests
|
|
try await client.withBatch { _ in
|
|
// No requests added
|
|
}
|
|
|
|
// Check that no messages were sent
|
|
#expect(await transport.sentMessages.isEmpty)
|
|
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Notify method sends notifications")
|
|
func testClientNotify() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Create a test notification
|
|
let notification = InitializedNotification.message()
|
|
try await client.notify(notification)
|
|
|
|
// Verify notification was sent
|
|
#expect(await transport.sentMessages.count == 1)
|
|
|
|
if let sentMessage = await transport.sentMessages.first,
|
|
let data = sentMessage.data(using: .utf8)
|
|
{
|
|
|
|
// Decode as Message<InitializedNotification>
|
|
let decoder = JSONDecoder()
|
|
do {
|
|
let decodedNotification = try decoder.decode(
|
|
Message<InitializedNotification>.self, from: data)
|
|
#expect(decodedNotification.method == InitializedNotification.name)
|
|
} catch {
|
|
#expect(Bool(false), "Failed to decode notification: \(error)")
|
|
}
|
|
} else {
|
|
#expect(Bool(false), "No message was sent")
|
|
}
|
|
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Initialize sends initialized notification")
|
|
func testClientInitializeNotification() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Create a task for initialize
|
|
let initTask = Task {
|
|
// Queue a response for the initialize request
|
|
try await Task.sleep(for: .milliseconds(10)) // Wait for request to be sent
|
|
|
|
if let lastMessage = await transport.sentMessages.last,
|
|
let data = lastMessage.data(using: .utf8),
|
|
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
|
{
|
|
|
|
// Create a valid initialize response
|
|
let response = Initialize.response(
|
|
id: request.id,
|
|
result: .init(
|
|
protocolVersion: Version.latest,
|
|
capabilities: .init(),
|
|
serverInfo: .init(name: "TestServer", version: "1.0"),
|
|
instructions: nil
|
|
)
|
|
)
|
|
|
|
try await transport.queue(response: response)
|
|
|
|
// Now complete the initialize call
|
|
_ = try await client.initialize()
|
|
|
|
// Verify that two messages were sent: initialize request and initialized notification
|
|
#expect(await transport.sentMessages.count == 2)
|
|
|
|
// Check that the second message is the initialized notification
|
|
let notifications = await transport.sentMessages
|
|
if notifications.count >= 2 {
|
|
let notificationJson = notifications[1]
|
|
if let notificationData = notificationJson.data(using: .utf8) {
|
|
do {
|
|
let decoder = JSONDecoder()
|
|
let decodedNotification = try decoder.decode(
|
|
Message<InitializedNotification>.self, from: notificationData)
|
|
#expect(decodedNotification.method == InitializedNotification.name)
|
|
} catch {
|
|
#expect(Bool(false), "Failed to decode notification: \(error)")
|
|
}
|
|
} else {
|
|
#expect(Bool(false), "Could not convert notification to data")
|
|
}
|
|
} else {
|
|
#expect(
|
|
Bool(false), "Expected both initialize request and initialized notification"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Wait with timeout
|
|
let timeoutTask = Task {
|
|
try await Task.sleep(for: .seconds(1))
|
|
initTask.cancel()
|
|
}
|
|
|
|
// Wait for the task to complete
|
|
do {
|
|
_ = try await initTask.value
|
|
} catch is CancellationError {
|
|
#expect(Bool(false), "Test timed out")
|
|
} catch {
|
|
#expect(Bool(false), "Unexpected error: \(error)")
|
|
}
|
|
|
|
timeoutTask.cancel()
|
|
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Race condition between send error and response")
|
|
func testSendErrorResponseRace() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Set up the transport to fail sends from the start
|
|
await transport.setFailSend(true)
|
|
|
|
// Create a ping request to get the ID
|
|
let request = Ping.request()
|
|
|
|
// Create a response for the request and queue it immediately
|
|
let response = Response<Ping>(id: request.id, result: .init())
|
|
let anyResponse = try AnyResponse(response)
|
|
try await transport.queue(response: anyResponse)
|
|
|
|
// Now attempt to send the request - this should fail due to send error
|
|
// but the response handler might also try to process the queued response
|
|
do {
|
|
_ = try await client.ping()
|
|
#expect(Bool(false), "Expected send to fail")
|
|
} catch let error as MCPError {
|
|
if case .transportError = error {
|
|
#expect(Bool(true))
|
|
} else {
|
|
#expect(Bool(false), "Expected transport error, got \(error)")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCPError, got \(error)")
|
|
}
|
|
|
|
// Verify no continuation misuse occurred
|
|
// (If it did, the test would have crashed)
|
|
|
|
await client.disconnect()
|
|
}
|
|
|
|
@Test("Race condition between response and send error")
|
|
func testResponseSendErrorRace() async throws {
|
|
let transport = MockTransport()
|
|
let client = Client(name: "TestClient", version: "1.0")
|
|
|
|
try await client.connect(transport: transport)
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
|
|
// Create a ping request to get the ID
|
|
let request = Ping.request()
|
|
|
|
// Create a response for the request and queue it immediately
|
|
let response = Response<Ping>(id: request.id, result: .init())
|
|
let anyResponse = try AnyResponse(response)
|
|
try await transport.queue(response: anyResponse)
|
|
|
|
// Set up the transport to fail sends
|
|
await transport.setFailSend(true)
|
|
|
|
// Now attempt to send the request
|
|
// The response might be processed before the send error occurs
|
|
do {
|
|
_ = try await client.ping()
|
|
// In this case, the response handler won the race and the request succeeded
|
|
#expect(Bool(true), "Response handler won the race - request succeeded")
|
|
} catch let error as MCPError {
|
|
if case .transportError = error {
|
|
// In this case, the send error handler won the race
|
|
#expect(Bool(true), "Send error handler won the race - request failed")
|
|
} else {
|
|
#expect(Bool(false), "Expected transport error, got \(error)")
|
|
}
|
|
} catch {
|
|
#expect(Bool(false), "Expected MCPError, got \(error)")
|
|
}
|
|
|
|
// Verify no continuation misuse occurred
|
|
// (If it did, the test would have crashed)
|
|
|
|
await client.disconnect()
|
|
}
|
|
}
|