modelcontextprotocol/swift-sdk
Watch
1
0
Fork
You've already forked swift-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update Client.connect(transport:) to automatically send initialize request (#100)

Browse source 
* Update Client to automatically call initialize in connect

* Make Client.connect(transport:) return value discardable

* Fix example usage in README

* Update example code in transports to remove initialize call

* Deprecate initialize method

* Update tests

* Update README

* Fix sampling tests

* Fix race condition tests
This commit is contained in:
Mattt 2025-05-26 06:04:38 -07:00 • committed by GitHub
commit a59b50bf5a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 302 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

337
Tests/MCPTests/ClientTests.swift
View file

@ -11,43 +11,33 @@ struct ClientTests {
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
// Set up a task to handle the initialize response
let initTask = Task {
try await client.initialize()
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
// 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
let result = try await client.connect(transport: transport)
#expect(await transport.isConnected == true)
#expect(result.protocolVersion == Version.latest)
await client.disconnect()
try await Task.sleep(for: .milliseconds(50))
#expect(await transport.isConnected == false)
initTask.cancel()
}
@Test(
@ -58,24 +48,50 @@ struct ClientTests {
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))
// Queue a response for the initialize request
try await Task.sleep(for: .milliseconds(10)) // Wait for request to be sent
// Create a task for the ping that we'll cancel
let pingTask = Task {
try await client.ping()
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 connect call which will automatically initialize
let result = try await client.connect(transport: transport)
#expect(result.protocolVersion == Version.latest)
#expect(result.serverInfo.name == "TestServer")
#expect(result.serverInfo.version == "1.0")
// Small delay to ensure message loop is started
try await Task.sleep(for: .milliseconds(10))
// Create a task for the ping
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 == 2) // Initialize + Ping
#expect(await transport.sentMessages.last?.contains(Ping.name) == true)
// Cancel the ping task
pingTask.cancel()
}
// 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))
@ -104,10 +120,35 @@ struct ClientTests {
@Test("Send failure handling")
func testClientSendFailure() async throws {
let transport = MockTransport()
await transport.setFailSend(true)
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
// Connect first without failure
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
initTask.cancel()
// Now set the transport to fail sends
await transport.setFailSend(true)
do {
try await client.ping()
@ -116,11 +157,13 @@ struct ClientTests {
if case MCPError.transportError = error {
#expect(Bool(true))
} else {
#expect(Bool(false), "Expected transport error")
#expect(Bool(false), "Expected transport error, got \(error)")
}
} catch {
#expect(Bool(false), "Expected MCP.Error")
}
await client.disconnect()
}
@Test("Strict configuration - capabilities check")
@ -129,6 +172,26 @@ struct ClientTests {
let config = Client.Configuration.strict
let client = Client(name: "TestClient", version: "1.0", configuration: config)
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
// Create a task for listPrompts
@ -152,6 +215,7 @@ struct ClientTests {
// Cancel the task if it's still running
promptsTask.cancel()
initTask.cancel()
// Disconnect client
await client.disconnect()
@ -164,8 +228,31 @@ struct ClientTests {
let config = Client.Configuration.default
let client = Client(name: "TestClient", version: "1.0", configuration: config)
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
// Make sure init task is complete
initTask.cancel()
// Wait a bit for any setup to complete
try await Task.sleep(for: .milliseconds(10))
@ -232,8 +319,29 @@ struct ClientTests {
let transport = MockTransport()
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10)) // Allow connection tasks
initTask.cancel()
let request1 = Ping.request()
let request2 = Ping.request()
@ -245,11 +353,11 @@ struct ClientTests {
resultTask2 = try await batch.addRequest(request2)
}
// Check if one batch message was sent
// Check if batch message was sent (after initialize and initialized notification)
let sentMessages = await transport.sentMessages
#expect(sentMessages.count == 1)
#expect(sentMessages.count == 3) // Initialize request + Initialized notification + Batch
guard let batchData = sentMessages.first?.data(using: .utf8) else {
guard let batchData = sentMessages.last?.data(using: .utf8) else {
#expect(Bool(false), "Failed to get batch data")
return
}
@ -291,8 +399,29 @@ struct ClientTests {
let transport = MockTransport()
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
initTask.cancel()
let request1 = Ping.request() // Success
let request2 = Ping.request() // Error
@ -304,8 +433,8 @@ struct ClientTests {
resultTasks.append(try await batch.addRequest(request2))
}
// Check if one batch message was sent
#expect(await transport.sentMessages.count == 1)
// Check if batch message was sent (after initialize and initialized notification)
#expect(await transport.sentMessages.count == 3) // Initialize request + Initialized notification + Batch
// Prepare batch response (success for 1, error for 2)
let response1 = Response<Ping>(id: request1.id, result: .init())
@ -350,16 +479,37 @@ struct ClientTests {
let transport = MockTransport()
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
initTask.cancel()
// 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)
// Check that only initialize message and initialized notification were sent
#expect(await transport.sentMessages.count == 2) // Initialize request + Initialized notification
await client.disconnect()
}
@ -369,17 +519,38 @@ struct ClientTests {
let transport = MockTransport()
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
initTask.cancel()
// Create a test notification
let notification = InitializedNotification.message()
try await client.notify(notification)
// Verify notification was sent
#expect(await transport.sentMessages.count == 1)
// Verify notification was sent (in addition to initialize and initialized notification)
#expect(await transport.sentMessages.count == 3) // Initialize request + Initialized notification + Custom notification
if let sentMessage = await transport.sentMessages.first,
if let sentMessage = await transport.sentMessages.last,
let data = sentMessage.data(using: .utf8)
{
@ -404,9 +575,6 @@ struct ClientTests {
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
@ -431,7 +599,8 @@ struct ClientTests {
try await transport.queue(response: response)
// Now complete the initialize call
_ = try await client.initialize()
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
// Verify that two messages were sent: initialize request and initialized notification
#expect(await transport.sentMessages.count == 2)
@ -485,8 +654,29 @@ struct ClientTests {
let transport = MockTransport()
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
initTask.cancel()
// Set up the transport to fail sends from the start
await transport.setFailSend(true)
@ -525,8 +715,29 @@ struct ClientTests {
let transport = MockTransport()
let client = Client(name: "TestClient", version: "1.0")
// Set up a task to handle the initialize response
let initTask = Task {
try await Task.sleep(for: .milliseconds(10))
if let lastMessage = await transport.sentMessages.last,
let data = lastMessage.data(using: .utf8),
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
{
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)
}
}
try await client.connect(transport: transport)
try await Task.sleep(for: .milliseconds(10))
initTask.cancel()
// Create a ping request to get the ID
let request = Ping.request()

5
Tests/MCPTests/HTTPClientTransportTests.swift
View file

@ -653,11 +653,8 @@ import Testing
}
}
// Execute the complete flow
try await client.connect(transport: transport)
// Step 1: Initialize client
let initResult = try await client.initialize()
let initResult = try await client.connect(transport: transport)
#expect(initResult.protocolVersion == Version.latest)
#expect(initResult.capabilities.tools != nil)

3
Tests/MCPTests/RoundtripTests.swift
View file

@ -99,10 +99,9 @@ struct RoundtripTests {
let client = Client(name: "TestClient", version: "1.0")
try await server.start(transport: serverTransport)
try await client.connect(transport: clientTransport)
let initTask = Task {
let result = try await client.initialize()
let result = try await client.connect(transport: clientTransport)
#expect(result.serverInfo.name == "TestServer")
#expect(result.serverInfo.version == "1.0.0")

24
Tests/MCPTests/SamplingTests.swift
View file

@ -587,30 +587,6 @@ struct SamplingIntegrationTests {
try await server.start(transport: serverTransport)
try await client.connect(transport: clientTransport)
// Test initialization and capability negotiation
let initTask = Task {
let result = try await client.initialize()
#expect(result.serverInfo.name == "SamplingTestServer")
#expect(result.serverInfo.version == "1.0.0")
#expect(
result.capabilities.sampling != nil, "Server should advertise sampling capability")
#expect(result.protocolVersion == Version.latest)
}
try await withThrowingTaskGroup(of: Void.self) { group in
group.addTask {
try await Task.sleep(for: .seconds(1))
initTask.cancel()
throw CancellationError()
}
group.addTask {
try await initTask.value
}
try await group.next()
group.cancelAll()
}
await server.stop()
await client.disconnect()
try? clientToServerRead.close()