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

README.md

@@ -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)
+let result = try await client.connect(transport: transport)
-
-// Initialize the connection
-let result = try await client.initialize()
 
 // 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",

Sources/MCP/Base/Transports/HTTPClientTransport.swift

@@ -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
 /// ```

Sources/MCP/Base/Transports/NetworkTransport.swift

@@ -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.

Sources/MCP/Base/Transports/StdioTransport.swift

@@ -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

Sources/MCP/Client/Client.swift

@@ -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,

Tests/MCPTests/ClientTests.swift

@@ -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(
+        // Set up a task to handle the initialize response
-        "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()
+            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
+        let result = try await client.connect(transport: transport)
-        try await Task.sleep(for: .milliseconds(10))
+        #expect(await transport.isConnected == true)
-
+        #expect(result.protocolVersion == Version.latest)
-        #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))
+        #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)
+        // Queue a response for the initialize request
-        // Small delay to ensure message loop is started
+        try await Task.sleep(for: .milliseconds(10))  // Wait for request to be sent
-        try await Task.sleep(for: .milliseconds(10))
 
-        // Create a task for the ping that we'll cancel
+        if let lastMessage = await transport.sentMessages.last,
-        let pingTask = Task {
+            let data = lastMessage.data(using: .utf8),
-            try await client.ping()
+            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
+        // Check if batch message was sent (after initialize and initialized notification)
-        #expect(await transport.sentMessages.count == 1)
+        #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
+        // Check that only initialize message and initialized notification were sent
-        #expect(await transport.sentMessages.isEmpty)
+        #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
+        // Verify notification was sent (in addition to initialize and initialized notification)
-        #expect(await transport.sentMessages.count == 1)
+        #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()

Tests/MCPTests/HTTPClientTransportTests.swift

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

Tests/MCPTests/RoundtripTests.swift

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

Tests/MCPTests/SamplingTests.swift

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