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

* 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
9 changed files with 302 additions and 106 deletions

View file

@ -54,10 +54,7 @@ let client = Client(name: "MyApp", version: "1.0.0")
// Create a transport and connect
let transport = StdioTransport()
try await client.connect(transport: transport)
// Initialize the connection
let result = try await client.initialize()
let result = try await client.connect(transport: transport)
// Check server capabilities
if result.capabilities.tools != nil {
@ -65,6 +62,11 @@ if result.capabilities.tools != nil {
}
```
> [!NOTE]
> The `Client.connect(transport:)` method returns the initialization result.
> This return value is discardable,
> so you can ignore it if you don't need to check server capabilities.
### Transport Options for Clients
#### Stdio Transport
@ -278,7 +280,7 @@ Handle common client errors:
```swift
do {
let result = try await client.initialize()
try await client.connect(transport: transport)
// Success
} catch let error as MCPError {
print("MCP Error: \(error.localizedDescription)")
@ -402,7 +404,7 @@ The server component allows your application to host model capabilities and resp
```swift
import MCP
// Initialize the server with capabilities
// Create a server with given capabilities
let server = Server(
name: "MyModelServer",
version: "1.0.0",

View file

@ -41,9 +41,6 @@ import Logging
/// let client = Client(name: "MyApp", version: "1.0.0")
/// try await client.connect(transport: transport)
///
/// // Initialize the connection
/// let result = try await client.initialize()
///
/// // The transport will automatically handle SSE events
/// // and deliver them through the client's notification handlers
/// ```

View file

@ -51,9 +51,6 @@ import Logging
/// // Use the transport with an MCP client
/// let client = Client(name: "MyApp", version: "1.0.0")
/// try await client.connect(transport: transport)
///
/// // Initialize the connection
/// let result = try await client.initialize()
/// ```
public actor NetworkTransport: Transport {
/// Represents a heartbeat message for connection health monitoring.

View file

@ -44,9 +44,6 @@ import struct Foundation.Data
/// // Create a transport and connect
/// let transport = StdioTransport()
/// try await client.connect(transport: transport)
///
/// // Initialize the connection
/// let result = try await client.initialize()
/// ```
public actor StdioTransport: Transport {
private let input: FileDescriptor

View file

@ -168,7 +168,8 @@ public actor Client {
}
/// Connect to the server using the given transport
public func connect(transport: any Transport) async throws {
@discardableResult
public func connect(transport: any Transport) async throws -> Initialize.Result {
self.connection = transport
try await self.connection?.connect()
@ -217,6 +218,9 @@ public actor Client {
} while true
await self.logger?.info("Client message handling loop task is terminating.")
}
// Automatically initialize after connecting
return try await _initialize()
}
/// Disconnect the client and cancel all pending requests
@ -480,7 +484,23 @@ public actor Client {
// MARK: - Lifecycle
/// Initialize the connection with the server.
///
/// - Important: This method is deprecated. Initialization now happens automatically
/// when calling `connect(transport:)`. You should use that method instead.
///
/// - Returns: The server's initialization response containing capabilities and server info
@available(
*, deprecated,
message:
"Initialization now happens automatically during connect. Use connect(transport:) instead."
)
public func initialize() async throws -> Initialize.Result {
return try await _initialize()
}
/// Internal initialization implementation
private func _initialize() async throws -> Initialize.Result {
let request = Initialize.request(
.init(
protocolVersion: Version.latest,

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()

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)

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")

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()