Prevent double continuation resumption in client (#113)

* 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
This commit is contained in:
Mattt 2025-05-26 04:36:01 -07:00 • committed by GitHub
commit fec680a243
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -479,4 +479,86 @@ struct ClientTests {
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()
}
}