Improve NetworkTransport and add test coverage (#108)

* Improve NetworkTransport and add test coverage

* Remove the weakTransport check since actors have special memory management rules

* Use Task.sleep(for:) consistently

* Always start connection in connect

* Run NetworkTransportTests serially

Sources/MCP/Base/Transports/NetworkTransport.swift

@@ -1,10 +1,29 @@
+import Foundation
 import Logging
 
-import struct Foundation.Data
-
 #if canImport(Network)
     import Network
 
+    /// Protocol that abstracts the Network.NWConnection functionality needed for NetworkTransport
+    @preconcurrency protocol NetworkConnectionProtocol {
+        var state: NWConnection.State { get }
+        var stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))? { get set }
+
+        func start(queue: DispatchQueue)
+        func cancel()
+        func send(
+            content: Data?, contentContext: NWConnection.ContentContext, isComplete: Bool,
+            completion: NWConnection.SendCompletion)
+        func receive(
+            minimumIncompleteLength: Int, maximumLength: Int,
+            completion: @escaping @Sendable (
+                Data?, NWConnection.ContentContext?, Bool, NWError?
+            ) -> Void)
+    }
+
+    /// Extension to conform NWConnection to internal NetworkConnectionProtocol
+    extension NWConnection: NetworkConnectionProtocol {}
+
     /// An implementation of a custom MCP transport using Apple's Network framework.
     ///
     /// This transport allows MCP clients and servers to communicate over TCP/UDP connections
@@ -37,23 +56,209 @@ import struct Foundation.Data
     /// let result = try await client.initialize()
     /// ```
     public actor NetworkTransport: Transport {
-        private let connection: NWConnection
+        /// Represents a heartbeat message for connection health monitoring.
-        /// Logger instance for transport-related events
+        public struct Heartbeat: RawRepresentable, Hashable, Sendable {
-        public nonisolated let logger: Logger
+            /// Magic bytes used to identify a heartbeat message.
+            private static let magicBytes: [UInt8] = [0xF0, 0x9F, 0x92, 0x93]
 
+            /// The timestamp of when the heartbeat was created.
+            public let timestamp: Date
+
+            /// Creates a new heartbeat with the current timestamp.
+            public init() {
+                self.timestamp = Date()
+            }
+
+            /// Creates a heartbeat with a specific timestamp.
+            ///
+            /// - Parameter timestamp: The timestamp for the heartbeat.
+            public init(timestamp: Date) {
+                self.timestamp = timestamp
+            }
+
+            // MARK: - RawRepresentable
+
+            public typealias RawValue = [UInt8]
+
+            /// Creates a heartbeat from its raw representation.
+            ///
+            /// - Parameter rawValue: The raw bytes of the heartbeat message.
+            /// - Returns: A heartbeat if the raw value is valid, nil otherwise.
+            public init?(rawValue: [UInt8]) {
+                // Check if the data has the correct format (magic bytes + timestamp)
+                guard rawValue.count >= 12,
+                    rawValue.prefix(4).elementsEqual(Self.magicBytes)
+                else {
+                    return nil
+                }
+
+                // Extract the timestamp
+                let timestampData = Data(rawValue[4..<12])
+                let timestamp = timestampData.withUnsafeBytes {
+                    $0.load(as: UInt64.self)
+                }
+
+                self.timestamp = Date(
+                    timeIntervalSinceReferenceDate: TimeInterval(timestamp) / 1000.0)
+            }
+
+            /// Converts the heartbeat to its raw representation.
+            public var rawValue: [UInt8] {
+                var result = Data(Self.magicBytes)
+
+                // Add timestamp (milliseconds since reference date)
+                let timestamp = UInt64(self.timestamp.timeIntervalSinceReferenceDate * 1000)
+                withUnsafeBytes(of: timestamp) { buffer in
+                    result.append(contentsOf: buffer)
+                }
+
+                return Array(result)
+            }
+
+            /// Converts the heartbeat to Data.
+            public var data: Data {
+                return Data(self.rawValue)
+            }
+
+            /// Checks if the given data represents a heartbeat message.
+            ///
+            /// - Parameter data: The data to check.
+            /// - Returns: true if the data is a heartbeat message, false otherwise.
+            public static func isHeartbeat(_ data: Data) -> Bool {
+                guard data.count >= 4 else {
+                    return false
+                }
+
+                return data.prefix(4).elementsEqual(Self.magicBytes)
+            }
+
+            /// Attempts to parse a heartbeat from the given data.
+            ///
+            /// - Parameter data: The data to parse.
+            /// - Returns: A heartbeat if the data is valid, nil otherwise.
+            public static func from(data: Data) -> Heartbeat? {
+                guard data.count >= 12 else {
+                    return nil
+                }
+
+                return Heartbeat(rawValue: Array(data))
+            }
+        }
+
+        /// Configuration for heartbeat behavior.
+        public struct HeartbeatConfiguration: Hashable, Sendable {
+            /// Whether heartbeats are enabled.
+            public let enabled: Bool
+            /// Interval between heartbeats in seconds.
+            public let interval: TimeInterval
+
+            /// Creates a new heartbeat configuration.
+            ///
+            /// - Parameters:
+            ///   - enabled: Whether heartbeats are enabled (default: true)
+            ///   - interval: Interval in seconds between heartbeats (default: 15.0)
+            public init(enabled: Bool = true, interval: TimeInterval = 15.0) {
+                self.enabled = enabled
+                self.interval = interval
+            }
+
+            /// Default heartbeat configuration.
+            public static let `default` = HeartbeatConfiguration()
+
+            /// Configuration with heartbeats disabled.
+            public static let disabled = HeartbeatConfiguration(enabled: false)
+        }
+
+        /// Configuration for connection retry behavior.
+        public struct ReconnectionConfiguration: Hashable, Sendable {
+            /// Whether the transport should attempt to reconnect on failure.
+            public let enabled: Bool
+            /// Maximum number of reconnection attempts.
+            public let maxAttempts: Int
+            /// Multiplier for exponential backoff on reconnect.
+            public let backoffMultiplier: Double
+
+            /// Creates a new reconnection configuration.
+            ///
+            /// - Parameters:
+            ///   - enabled: Whether reconnection should be attempted on failure (default: true)
+            ///   - maxAttempts: Maximum number of reconnection attempts (default: 5)
+            ///   - backoffMultiplier: Multiplier for exponential backoff on reconnect (default: 1.5)
+            public init(
+                enabled: Bool = true,
+                maxAttempts: Int = 5,
+                backoffMultiplier: Double = 1.5
+            ) {
+                self.enabled = enabled
+                self.maxAttempts = maxAttempts
+                self.backoffMultiplier = backoffMultiplier
+            }
+
+            /// Default reconnection configuration.
+            public static let `default` = ReconnectionConfiguration()
+
+            /// Configuration with reconnection disabled.
+            public static let disabled = ReconnectionConfiguration(enabled: false)
+
+            /// Calculates the backoff delay for a given attempt number.
+            ///
+            /// - Parameter attempt: The current attempt number (1-based)
+            /// - Returns: The delay in seconds before the next attempt
+            public func backoffDelay(for attempt: Int) -> TimeInterval {
+                let baseDelay = 0.5  // 500ms
+                return baseDelay * pow(backoffMultiplier, Double(attempt - 1))
+            }
+        }
+
+        // State tracking
         private var isConnected = false
+        private var isStopping = false
+        private var reconnectAttempt = 0
+        private var heartbeatTask: Task<Void, Never>?
+        private var lastHeartbeatTime: Date?
         private let messageStream: AsyncThrowingStream<Data, Swift.Error>
         private let messageContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation
 
         // Track connection state for continuations
         private var connectionContinuationResumed = false
 
+        // Connection is marked nonisolated(unsafe) to allow access from closures
+        private nonisolated(unsafe) var connection: NetworkConnectionProtocol
+
+        /// Logger instance for transport-related events
+        public nonisolated let logger: Logger
+
+        // Configuration
+        private let heartbeatConfig: HeartbeatConfiguration
+        private let reconnectionConfig: ReconnectionConfiguration
+
         /// Creates a new NetworkTransport with the specified NWConnection
         ///
         /// - Parameters:
         ///   - connection: The NWConnection to use for communication
         ///   - logger: Optional logger instance for transport events
-        public init(connection: NWConnection, logger: Logger? = nil) {
+        ///   - reconnectionConfig: Configuration for reconnection behavior (default: .default)
+        ///   - heartbeatConfig: Configuration for heartbeat behavior (default: .default)
+        public init(
+            connection: NWConnection,
+            logger: Logger? = nil,
+            heartbeatConfig: HeartbeatConfiguration = .default,
+            reconnectionConfig: ReconnectionConfiguration = .default
+        ) {
+            self.init(
+                connection,
+                logger: logger,
+                heartbeatConfig: heartbeatConfig,
+                reconnectionConfig: reconnectionConfig
+            )
+        }
+
+        init(
+            _ connection: NetworkConnectionProtocol,
+            logger: Logger? = nil,
+            heartbeatConfig: HeartbeatConfiguration = .default,
+            reconnectionConfig: ReconnectionConfiguration = .default
+        ) {
             self.connection = connection
             self.logger =
                 logger
@@ -61,6 +266,8 @@ import struct Foundation.Data
                     label: "mcp.transport.network",
                     factory: { _ in SwiftLogNoOpLogHandler() }
                 )
+            self.reconnectionConfig = reconnectionConfig
+            self.heartbeatConfig = heartbeatConfig
 
             // Create message stream
             var continuation: AsyncThrowingStream<Data, Swift.Error>.Continuation!
@@ -77,6 +284,10 @@ import struct Foundation.Data
         public func connect() async throws {
             guard !isConnected else { return }
 
+            // Reset state for fresh connection
+            isStopping = false
+            reconnectAttempt = 0
+
             // Reset continuation state
             connectionContinuationResumed = false
 
@@ -100,21 +311,19 @@ import struct Foundation.Data
                                 error: error, continuation: continuation)
                         case .cancelled:
                             await self.handleConnectionCancelled(continuation: continuation)
-                        default:
+                        case .waiting(let error):
-                            // Wait for ready or failed state
+                            self.logger.debug("Connection waiting: \(error)")
-                            break
+                        case .preparing:
+                            self.logger.debug("Connection preparing...")
+                        case .setup:
+                            self.logger.debug("Connection setup...")
+                        @unknown default:
+                            self.logger.warning("Unknown connection state")
                         }
                     }
                 }
 
-                // Start the connection if it's not already started
+                connection.start(queue: .main)
-                if connection.state != .ready {
-                    connection.start(queue: .main)
-                } else {
-                    Task { @MainActor in
-                        await self.handleConnectionReady(continuation: continuation)
-                    }
-                }
             }
         }
 
@@ -127,13 +336,84 @@ import struct Foundation.Data
             if !connectionContinuationResumed {
                 connectionContinuationResumed = true
                 isConnected = true
+
+                // Reset reconnect attempt counter on successful connection
+                reconnectAttempt = 0
                 logger.info("Network transport connected successfully")
                 continuation.resume()
+
                 // Start the receive loop after connection is established
                 Task { await self.receiveLoop() }
+
+                // Start heartbeat task if enabled
+                if heartbeatConfig.enabled {
+                    startHeartbeat()
+                }
             }
         }
 
+        /// Starts a task to periodically send heartbeats to check connection health
+        private func startHeartbeat() {
+            // Cancel any existing heartbeat task
+            heartbeatTask?.cancel()
+
+            // Start a new heartbeat task
+            heartbeatTask = Task { [weak self] in
+                guard let self = self else { return }
+
+                // Initial delay before starting heartbeats
+                try? await Task.sleep(for: .seconds(1))
+
+                while !Task.isCancelled {
+                    do {
+                        // Check actor-isolated properties first
+                        let isStopping = await self.isStopping
+                        let isConnected = await self.isConnected
+
+                        guard !isStopping && isConnected else { break }
+
+                        try await self.sendHeartbeat()
+                        try await Task.sleep(for: .seconds(self.heartbeatConfig.interval))
+                    } catch {
+                        // If heartbeat fails, log and retry after a shorter interval
+                        self.logger.warning("Heartbeat failed: \(error)")
+                        try? await Task.sleep(for: .seconds(2))
+                    }
+                }
+            }
+        }
+
+        /// Sends a heartbeat message to verify connection health
+        private func sendHeartbeat() async throws {
+            guard isConnected && !isStopping else { return }
+
+            // Try to send the heartbeat (without the newline delimiter used for normal messages)
+            try await withCheckedThrowingContinuation {
+                [weak self] (continuation: CheckedContinuation<Void, Swift.Error>) in
+                guard let self = self else {
+                    continuation.resume(throwing: MCPError.internalError("Transport deallocated"))
+                    return
+                }
+
+                connection.send(
+                    content: Heartbeat().data,
+                    contentContext: .defaultMessage,
+                    isComplete: true,
+                    completion: .contentProcessed { [weak self] error in
+                        if let error = error {
+                            continuation.resume(throwing: error)
+                        } else {
+                            Task { [weak self] in
+                                await self?.setLastHeartbeatTime(Date())
+                            }
+                            continuation.resume()
+                        }
+                    })
+            }
+
+            logger.debug("Heartbeat sent")
+        }
+
         /// Handles connection failure
         ///
         /// - Parameters:
@@ -145,7 +425,12 @@ import struct Foundation.Data
             if !connectionContinuationResumed {
                 connectionContinuationResumed = true
                 logger.error("Connection failed: \(error)")
-                continuation.resume(throwing: error)
+
+                await handleReconnection(
+                    error: error,
+                    continuation: continuation,
+                    context: "failure"
+                )
             }
         }
 
@@ -158,7 +443,55 @@ import struct Foundation.Data
             if !connectionContinuationResumed {
                 connectionContinuationResumed = true
                 logger.warning("Connection cancelled")
-                continuation.resume(throwing: MCPError.internalError("Connection cancelled"))
+
+                await handleReconnection(
+                    error: MCPError.internalError("Connection cancelled"),
+                    continuation: continuation,
+                    context: "cancellation"
+                )
+            }
+        }
+
+        /// Common reconnection handling logic
+        ///
+        /// - Parameters:
+        ///   - error: The error that triggered the reconnection
+        ///   - continuation: The continuation to resume with the error
+        ///   - context: The context of the reconnection (for logging)
+        private func handleReconnection(
+            error: Swift.Error,
+            continuation: CheckedContinuation<Void, Swift.Error>,
+            context: String
+        ) async {
+            if !isStopping,
+                reconnectionConfig.enabled,
+                reconnectAttempt < reconnectionConfig.maxAttempts
+            {
+                // Try to reconnect with exponential backoff
+                reconnectAttempt += 1
+                logger.info(
+                    "Attempting reconnection after \(context) (\(reconnectAttempt)/\(reconnectionConfig.maxAttempts))..."
+                )
+
+                // Calculate backoff delay
+                let delay = reconnectionConfig.backoffDelay(for: reconnectAttempt)
+
+                // Schedule reconnection attempt after delay
+                Task {
+                    try? await Task.sleep(for: .seconds(delay))
+                    if !isStopping {
+                        // Cancel the current connection before attempting to reconnect.
+                        self.connection.cancel()
+                        // Resume original continuation with error; outer logic or a new call to connect() will handle retry.
+                        continuation.resume(throwing: error)
+                    } else {
+                        continuation.resume(throwing: error)  // Stopping, so fail.
+                    }
+                }
+            } else {
+                // Not configured to reconnect, exceeded max attempts, or stopping
+                self.connection.cancel()  // Ensure connection is cancelled
+                continuation.resume(throwing: error)
             }
         }
 
@@ -168,7 +501,15 @@ import struct Foundation.Data
         /// and releases associated resources.
         public func disconnect() async {
             guard isConnected else { return }
+
+            // Mark as stopping to prevent reconnection attempts during disconnect
+            isStopping = true
             isConnected = false
+
+            // Cancel heartbeat task if it exists
+            heartbeatTask?.cancel()
+            heartbeatTask = nil
+
             connection.cancel()
             messageContinuation.finish()
             logger.info("Network transport disconnected")
@@ -202,6 +543,8 @@ import struct Foundation.Data
 
                 connection.send(
                     content: messageWithNewline,
+                    contentContext: .defaultMessage,
+                    isComplete: true,
                     completion: .contentProcessed { [weak self] error in
                         guard let self = self else { return }
 
@@ -210,6 +553,36 @@ import struct Foundation.Data
                                 sendContinuationResumed = true
                                 if let error = error {
                                     self.logger.error("Send error: \(error)")
+
+                                    // Check if we should attempt to reconnect on send failure
+                                    let isStopping = await self.isStopping  // Await actor-isolated property
+                                    if !isStopping && self.reconnectionConfig.enabled {
+                                        let isConnected = await self.isConnected
+                                        if isConnected {
+                                            if error.isConnectionLost {
+                                                self.logger.warning(
+                                                    "Connection appears broken, will attempt to reconnect..."
+                                                )
+
+                                                // Schedule connection restart
+                                                Task { [weak self] in  // Operate on self's executor
+                                                    guard let self = self else { return }
+
+                                                    await self.setIsConnected(false)
+
+                                                    try? await Task.sleep(for: .milliseconds(500))
+
+                                                    let currentIsStopping = await self.isStopping
+                                                    if !currentIsStopping {
+                                                        // Cancel the connection, then attempt to reconnect fully.
+                                                        self.connection.cancel()
+                                                        try? await self.connect()
+                                                    }
+                                                }
+                                            }
+                                        }
+                                    }
+
                                     continuation.resume(
                                         throwing: MCPError.internalError("Send error: \(error)"))
                                 } else {
@@ -238,15 +611,50 @@ import struct Foundation.Data
         /// Messages are delimited by newline characters.
         private func receiveLoop() async {
             var buffer = Data()
+            var consecutiveEmptyReads = 0
+            let maxConsecutiveEmptyReads = 5
 
-            while isConnected && !Task.isCancelled {
+            while isConnected && !Task.isCancelled && !isStopping {
                 do {
                     let newData = try await receiveData()
-                    // Check for EOF (empty data)
+
+                    // Check for EOF or empty data
                     if newData.isEmpty {
-                        logger.info("Connection closed by peer (EOF).")
+                        consecutiveEmptyReads += 1
-                        break  // Exit loop gracefully
+
+                        if consecutiveEmptyReads >= maxConsecutiveEmptyReads {
+                            logger.warning(
+                                "Multiple consecutive empty reads (\(consecutiveEmptyReads)), possible connection issue"
+                            )
+
+                            // Check connection state
+                            if connection.state != .ready {
+                                logger.info("Connection no longer ready, exiting receive loop")
+                                break
+                            }
+                        }
+
+                        // Brief pause before retry
+                        try await Task.sleep(for: .milliseconds(100))
+                        continue
                     }
+
+                    // Check if this is a heartbeat message
+                    if Heartbeat.isHeartbeat(newData) {
+                        logger.debug("Received heartbeat from peer")
+
+                        // Extract timestamp if available
+                        if let heartbeat = Heartbeat.from(data: newData) {
+                            logger.debug("Heartbeat timestamp: \(heartbeat.timestamp)")
+                        }
+
+                        // Reset the counter since we got valid data
+                        consecutiveEmptyReads = 0
+                        continue  // Skip regular message processing for heartbeats
+                    }
+
+                    // Reset counter on successful data read
+                    consecutiveEmptyReads = 0
                     buffer.append(newData)
 
                     // Process complete messages
@@ -261,21 +669,97 @@ import struct Foundation.Data
                         }
                     }
                 } catch let error as NWError {
-                    if !Task.isCancelled {
+                    if !Task.isCancelled && !isStopping {
                         logger.error("Network error occurred", metadata: ["error": "\(error)"])
-                        messageContinuation.finish(throwing: MCPError.transportError(error))
+
+                        // Check for specific connection-related errors
+                        if error.isConnectionLost {
+                            // If we should reconnect, don't finish the message stream yet
+                            if reconnectionConfig.enabled
+                                && reconnectAttempt < reconnectionConfig.maxAttempts
+                            {
+                                reconnectAttempt += 1
+                                logger.info(
+                                    "Network connection lost, attempting reconnection (\(reconnectAttempt)/\(reconnectionConfig.maxAttempts))..."
+                                )
+
+                                // Mark as not connected while attempting reconnection
+                                isConnected = false
+
+                                // Schedule reconnection attempt
+                                Task {
+                                    let delay = reconnectionConfig.backoffDelay(
+                                        for: reconnectAttempt)
+                                    try? await Task.sleep(for: .seconds(delay))
+
+                                    if !isStopping {
+                                        // Cancel the connection, then attempt to reconnect fully.
+                                        self.connection.cancel()
+                                        try? await self.connect()
+
+                                        // If connect succeeded, a new receive loop will be started
+                                    }
+                                }
+
+                                // Exit this receive loop since we're starting a new one after reconnect
+                                break
+                            } else {
+                                // We're not reconnecting, finish the message stream with error
+                                messageContinuation.finish(
+                                    throwing: MCPError.transportError(error))
+                                break
+                            }
+                        } else {
+                            // For other network errors, log but continue trying
+                            do {
+                                try await Task.sleep(for: .milliseconds(100))  // 100ms pause
+                                continue
+                            } catch {
+                                logger.error("Failed to sleep after network error: \(error)")
+                                break
+                            }
+                        }
                     }
                     break
                 } catch {
-                    if !Task.isCancelled {
+                    if !Task.isCancelled && !isStopping {
                         logger.error("Receive error: \(error)")
-                        messageContinuation.finish(throwing: error)
+
+                        if reconnectionConfig.enabled
+                            && reconnectAttempt < reconnectionConfig.maxAttempts
+                        {
+                            // Similar reconnection logic for other errors
+                            reconnectAttempt += 1
+                            logger.info(
+                                "Error during receive, attempting reconnection (\(reconnectAttempt)/\(reconnectionConfig.maxAttempts))..."
+                            )
+
+                            isConnected = false
+
+                            Task {
+                                let delay = reconnectionConfig.backoffDelay(for: reconnectAttempt)
+                                try? await Task.sleep(for: .seconds(delay))
+
+                                if !isStopping {
+                                    self.connection.cancel()
+                                    try? await connect()
+                                }
+                            }
+
+                            break
+                        } else {
+                            messageContinuation.finish(throwing: error)
+                        }
                     }
                     break
                 }
             }
 
-            messageContinuation.finish()
+            // If stopping normally, finish the stream without error
+            if isStopping {
+                logger.debug("Receive loop stopping normally")
+                messageContinuation.finish()
+            }
         }
 
         /// Receives a chunk of data from the network connection
@@ -293,7 +777,7 @@ import struct Foundation.Data
                 }
 
                 connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) {
-                    content, _, _, error in
+                    content, _, isComplete, error in
                     Task { @MainActor in
                         if !receiveContinuationResumed {
                             receiveContinuationResumed = true
@@ -301,6 +785,9 @@ import struct Foundation.Data
                                 continuation.resume(throwing: MCPError.transportError(error))
                             } else if let content = content {
                                 continuation.resume(returning: content)
+                            } else if isComplete {
+                                self.logger.debug("Connection completed by peer")
+                                continuation.resume(throwing: MCPError.connectionClosed)
                             } else {
                                 // EOF: Resume with empty data instead of throwing an error
                                 continuation.resume(returning: Data())
@@ -310,5 +797,22 @@ import struct Foundation.Data
                 }
             }
         }
+
+        private func setLastHeartbeatTime(_ time: Date) {
+            self.lastHeartbeatTime = time
+        }
+
+        private func setIsConnected(_ connected: Bool) {
+            self.isConnected = connected
+        }
+    }
+
+    extension NWError {
+        /// Whether this error indicates a connection has been lost or reset.
+        fileprivate var isConnectionLost: Bool {
+            let nsError = self as NSError
+            return nsError.code == 57  // Socket is not connected (EHOSTUNREACH or ENOTCONN)
+                || nsError.code == 54  // Connection reset by peer (ECONNRESET)
+        }
     }
 #endif

Tests/MCPTests/NetworkTransportTests.swift

@@ -0,0 +1,688 @@
+import Foundation
+import Logging
+import Testing
+
+@testable import MCP
+
+#if canImport(Network)
+    import Network
+
+    /// A mock implementation of NetworkConnectionProtocol for testing
+    final class MockNetworkConnection: NetworkConnectionProtocol, @unchecked Sendable {
+        /// Current state of the connection
+        private var mockState: NWConnection.State = .setup
+
+        /// Error to be returned on send/receive operations
+        private var mockError: Swift.Error?
+
+        /// Data queue for testing
+        private var dataToReceive: [Data] = []
+        private var sentData: [Data] = []
+
+        /// The state update handler
+        public var stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))?
+
+        /// Current state
+        var state: NWConnection.State {
+            return mockState
+        }
+
+        /// Initialize a mock connection
+        init() {}
+
+        /// Start the connection
+        func start(queue: DispatchQueue) {
+            // Simulate successful connection by default
+            Task { @MainActor in
+                self.updateState(.ready)
+            }
+        }
+
+        /// Send data through the connection
+        func send(
+            content: Data?,
+            contentContext: NWConnection.ContentContext,
+            isComplete: Bool,
+            completion: NWConnection.SendCompletion
+        ) {
+            if let content = content {
+                sentData.append(content)
+            }
+
+            switch completion {
+            case .contentProcessed(let handler):
+                Task { @MainActor in
+                    handler(self.mockError as? NWError)
+                }
+            default:
+                break
+            }
+        }
+
+        /// Receive data from the connection
+        func receive(
+            minimumIncompleteLength: Int,
+            maximumLength: Int,
+            completion: @escaping @Sendable (
+                Data?, NWConnection.ContentContext?, Bool, NWError?
+            ) -> Void
+        ) {
+            Task { @MainActor in
+                if self.mockState == .cancelled {
+                    completion(
+                        nil, nil, true,
+                        NWError.posix(POSIXErrorCode.ECANCELED)
+                    )
+                    return
+                }
+
+                if let error = self.mockError {
+                    completion(nil, nil, false, error as? NWError)
+                    return
+                }
+
+                if self.dataToReceive.isEmpty {
+                    completion(Data(), nil, false, nil)
+                    return
+                }
+
+                let data = self.dataToReceive.removeFirst()
+                completion(data, nil, self.dataToReceive.isEmpty, nil)
+            }
+        }
+
+        /// Cancel the connection
+        func cancel() {
+            updateState(.cancelled)
+        }
+
+        // Test helpers
+
+        /// Simulate the connection becoming ready
+        func simulateReady() {
+            updateState(.ready)
+        }
+
+        /// Simulate the connection becoming preparing
+        func simulatePreparing() {
+            updateState(.preparing)
+        }
+
+        /// Simulate a connection failure
+        func simulateFailure(
+            error: Swift.Error? = nil
+        ) {
+            mockError = error
+            if let nwError = error as? NWError {
+                updateState(.failed(nwError))
+            } else {
+                updateState(.failed(NWError.posix(POSIXErrorCode(rawValue: 57)!)))
+            }
+        }
+
+        /// Simulate connection cancellation
+        func simulateCancellation() {
+            updateState(.cancelled)
+        }
+
+        /// Update the connection state and notify handler
+        private func updateState(_ newState: NWConnection.State) {
+            mockState = newState
+            Task { @MainActor in
+                self.stateUpdateHandler?(newState)
+            }
+        }
+
+        /// Queue data to be received
+        func queueDataForReceiving(_ data: Data) {
+            dataToReceive.append(data)
+        }
+
+        /// Queue a heartbeat message to be received
+        func queueHeartbeat() {
+            // Create a mock heartbeat message that matches the format
+            let magicBytes: [UInt8] = [0xF0, 0x9F, 0x92, 0x93]  // Magic bytes for heartbeat
+            var data = Data(magicBytes)
+            let timestamp = UInt64(Date().timeIntervalSinceReferenceDate * 1000)
+            withUnsafeBytes(of: timestamp) { buffer in
+                data.append(contentsOf: buffer)
+            }
+            queueDataForReceiving(data)
+        }
+
+        /// Queue text message to be received
+        func queueTextMessage(_ text: String) {
+            guard let data = text.data(using: .utf8) else { return }
+            queueDataForReceiving(data)
+        }
+
+        /// Get all sent data
+        func getSentData() -> [Data] {
+            return sentData
+        }
+
+        /// Clear sent data buffer
+        func clearSentData() {
+            sentData.removeAll()
+        }
+    }
+
+    @Suite("Network Transport Tests", .serialized)
+    struct NetworkTransportTests {
+        @Test("Heartbeat Creation And Parsing")
+        func testHeartbeatCreationAndParsing() {
+            // Create a heartbeat
+            let heartbeat = NetworkTransport.Heartbeat()
+
+            // Convert to data and back
+            let data = heartbeat.data
+            let parsed = NetworkTransport.Heartbeat.from(data: data)
+
+            #expect(parsed != nil)
+
+            // Time should be very close (within 1 second)
+            if let parsed = parsed {
+                let timeDifference = abs(parsed.timestamp.timeIntervalSince(heartbeat.timestamp))
+                #expect(timeDifference < 1.0)
+            }
+
+            // Test invalid data
+            let invalidData = Data([0x01, 0x02, 0x03])
+            #expect(NetworkTransport.Heartbeat.from(data: invalidData) == nil)
+            #expect(NetworkTransport.Heartbeat.isHeartbeat(invalidData) == false)
+            #expect(NetworkTransport.Heartbeat.isHeartbeat(data) == true)
+        }
+
+        @Test("Reconnection Configuration")
+        func testReconnectionConfiguration() {
+            // Create custom config
+            let config = NetworkTransport.ReconnectionConfiguration(
+                enabled: true,
+                maxAttempts: 3,
+                backoffMultiplier: 2.0
+            )
+
+            #expect(config.enabled == true)
+            #expect(config.maxAttempts == 3)
+            #expect(config.backoffMultiplier == 2.0)
+
+            // Test backoff delay calculation
+            let firstDelay = config.backoffDelay(for: 1)
+            let secondDelay = config.backoffDelay(for: 2)
+            let thirdDelay = config.backoffDelay(for: 3)
+
+            // Check delays are approximately correct (within 0.001)
+            #expect(abs(firstDelay - 0.5) < 0.001)
+            #expect(abs(secondDelay - 1.0) < 0.001)
+            #expect(abs(thirdDelay - 2.0) < 0.001)
+
+            // Test disabled config
+            let disabledConfig = NetworkTransport.ReconnectionConfiguration.disabled
+            #expect(disabledConfig.enabled == false)
+        }
+
+        @Test("Heartbeat Configuration")
+        func testHeartbeatConfiguration() {
+            // Create custom config
+            let config = NetworkTransport.HeartbeatConfiguration(
+                enabled: true,
+                interval: 5.0
+            )
+
+            #expect(config.enabled == true)
+            #expect(config.interval == 5.0)
+
+            // Test default config
+            let defaultConfig = NetworkTransport.HeartbeatConfiguration.default
+            #expect(defaultConfig.enabled == true)
+            #expect(defaultConfig.interval == 15.0)
+
+            // Test disabled config
+            let disabledConfig = NetworkTransport.HeartbeatConfiguration.disabled
+            #expect(disabledConfig.enabled == false)
+        }
+
+        @Test("Connect Success")
+        func testNetworkTransportConnectSuccess() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled  // Disable heartbeats for simplified testing
+            )
+
+            try await transport.connect()
+
+            // Verify connection state
+            #expect(mockConnection.state == .ready)
+
+            await transport.disconnect()
+        }
+
+        @Test("Connect Failure")
+        func testNetworkTransportConnectFailure() async {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                reconnectionConfig: .disabled  // Disable reconnection for this test
+            )
+
+            // Simulate failure before connecting
+            mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
+
+            do {
+                try await transport.connect()
+                Issue.record("Expected connect to throw an error")
+            } catch let error as MCPError {
+                // Expected failure
+                #expect(error.localizedDescription.contains("Connection failed"))
+            } catch let error as NWError {
+                // Also accept NWError since it's the underlying error
+                #expect((error as NSError).code == POSIXErrorCode.ECONNRESET.rawValue)
+            } catch {
+                Issue.record("Unexpected error type: \(type(of: error))")
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Send Message")
+        func testNetworkTransportSendMessage() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            try await transport.connect()
+
+            // Test sending a simple message
+            let message = #"{"key":"value"}"#.data(using: .utf8)!
+            try await transport.send(message)
+
+            // Verify the message was sent with a newline delimiter
+            let sentData = mockConnection.getSentData()
+            #expect(sentData.count == 1)
+
+            if sentData.count > 0 {
+                let expectedOutput = message + "\n".data(using: .utf8)!
+                #expect(sentData[0] == expectedOutput)
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Receive Message")
+        func testNetworkTransportReceiveMessage() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            // Queue a message to be received
+            let message = #"{"key":"value"}"#
+            let messageWithNewline = message + "\n"
+            mockConnection.queueTextMessage(messageWithNewline)
+
+            try await transport.connect()
+
+            // Start receiving messages
+            let stream = await transport.receive()
+            var iterator = stream.makeAsyncIterator()
+
+            // Get first message
+            let received = try await iterator.next()
+            #expect(received != nil)
+
+            if let received = received {
+                #expect(received == message.data(using: .utf8)!)
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Heartbeat Send and Receive")
+        func testNetworkTransportHeartbeat() async throws {
+            let mockConnection = MockNetworkConnection()
+
+            // Create transport with rapid heartbeats
+            let heartbeatConfig = NetworkTransport.HeartbeatConfiguration(
+                enabled: true,
+                interval: 0.1  // Short interval for testing
+            )
+
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: heartbeatConfig
+            )
+
+            try await transport.connect()
+
+            // Wait for initial connection setup
+            try await Task.sleep(for: .milliseconds(100))
+
+            // Wait for initial heartbeat delay (1 second) plus a small buffer
+            try await Task.sleep(for: .seconds(1.2))
+
+            // Check if heartbeat was sent
+            let sentData = mockConnection.getSentData()
+            #expect(sentData.count >= 1, "No heartbeat was sent after \(sentData.count) attempts")
+
+            if let firstSent = sentData.first {
+                #expect(
+                    NetworkTransport.Heartbeat.isHeartbeat(firstSent),
+                    "Sent data is not a heartbeat")
+            }
+
+            // Queue a heartbeat to be received
+            mockConnection.queueHeartbeat()
+
+            // Wait for heartbeat processing
+            try await Task.sleep(for: .milliseconds(100))
+
+            await transport.disconnect()
+        }
+
+        @Test("Reconnection")
+        func testNetworkTransportReconnection() async throws {
+            let mockConnection = MockNetworkConnection()
+
+            // Configure for quick reconnection
+            let reconnectionConfig = NetworkTransport.ReconnectionConfiguration(
+                enabled: true,
+                maxAttempts: 2,
+                backoffMultiplier: 1.0
+            )
+
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled,
+                reconnectionConfig: reconnectionConfig
+            )
+
+            try await transport.connect()
+
+            // Simulate connection failure during operation
+            mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
+
+            // Wait a bit to ensure failure is processed
+            try await Task.sleep(for: .milliseconds(100))
+
+            // Try to send message after failure - should trigger reconnection process
+            let message = #"{"test":"reconnect"}"#.data(using: .utf8)!
+
+            do {
+                try await transport.send(message)
+                Issue.record("Expected send to fail after connection lost")
+            } catch {
+                // Expected error
+                #expect(error is MCPError, "Expected MCPError but got \(type(of: error))")
+            }
+
+            // Wait for potential reconnection attempt
+            try await Task.sleep(for: .milliseconds(600))
+
+            await transport.disconnect()
+        }
+
+        @Test("Multiple Messages")
+        func testNetworkTransportMultipleMessages() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            // Queue multiple messages
+            let messages = [
+                #"{"id":1,"method":"test1"}"#,
+                #"{"id":2,"method":"test2"}"#,
+                #"{"id":3,"method":"test3"}"#,
+            ]
+
+            for message in messages {
+                mockConnection.queueTextMessage(message + "\n")
+            }
+
+            try await transport.connect()
+
+            // Receive and verify all messages
+            let stream = await transport.receive()
+            var receiveCount = 0
+
+            for try await data in stream {
+                if let receivedStr = String(data: data, encoding: .utf8) {
+                    #expect(messages.contains(receivedStr))
+                    receiveCount += 1
+
+                    if receiveCount >= messages.count {
+                        break
+                    }
+                }
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Disconnect During Receive")
+        func testNetworkTransportDisconnectDuringReceive() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            try await transport.connect()
+
+            // Start a task to receive messages
+            let receiveTask = Task {
+                var count = 0
+                for try await _ in await transport.receive() {
+                    count += 1
+                    if count > 10 {
+                        // Prevent infinite loop in test
+                        break
+                    }
+                }
+            }
+
+            // Let the receive loop start
+            try await Task.sleep(for: .milliseconds(100))
+
+            // Disconnect while receiving
+            await transport.disconnect()
+
+            // Wait for the receive task to complete
+            _ = await receiveTask.result
+        }
+
+        @Test("Connection State Transitions")
+        func testConnectionStateTransitions() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            // Test setup -> preparing -> ready transition
+            mockConnection.simulatePreparing()
+            try await Task.sleep(for: .milliseconds(100))
+            mockConnection.simulateReady()
+            try await transport.connect()
+            #expect(mockConnection.state == .ready)
+
+            // Test ready -> failed transition
+            mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
+            try await Task.sleep(for: .milliseconds(100))
+            if case .failed = mockConnection.state {
+                // expected
+            } else {
+                Issue.record("Expected state to be failed")
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Partial Message Reception")
+        func testPartialMessageReception() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            try await transport.connect()
+
+            // Split a message into multiple parts
+            let message = #"{"key":"value"}"#
+            let parts = [
+                message.prefix(5).data(using: .utf8)!,
+                message.dropFirst(5).data(using: .utf8)!,
+                "\n".data(using: .utf8)!,
+            ]
+
+            // Queue the parts
+            for part in parts {
+                mockConnection.queueDataForReceiving(part)
+            }
+
+            // Start receiving messages
+            let stream = await transport.receive()
+            var iterator = stream.makeAsyncIterator()
+
+            // Get the complete message
+            let received = try await iterator.next()
+            #expect(received != nil)
+            if let received = received {
+                #expect(received == message.data(using: .utf8)!)
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Large Message Handling")
+        func testLargeMessageHandling() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled
+            )
+
+            try await transport.connect()
+
+            // Create a large message (larger than typical receive buffer)
+            let largeMessage = String(repeating: "x", count: 100_000)
+            let messageWithNewline = largeMessage + "\n"
+            mockConnection.queueTextMessage(messageWithNewline)
+
+            // Start receiving messages
+            let stream = await transport.receive()
+            var iterator = stream.makeAsyncIterator()
+
+            // Get the message
+            let received = try await iterator.next()
+            #expect(received != nil)
+            if let received = received {
+                #expect(received.count == largeMessage.count)
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Reconnection Backoff")
+        func testReconnectionBackoff() async throws {
+            let mockConnection = MockNetworkConnection()
+            let startTime = Date()
+
+            // Configure for quick reconnection with known backoff
+            let reconnectionConfig = NetworkTransport.ReconnectionConfiguration(
+                enabled: true,
+                maxAttempts: 3,
+                backoffMultiplier: 2.0
+            )
+
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled,
+                reconnectionConfig: reconnectionConfig
+            )
+
+            try await transport.connect()
+
+            // Simulate failure
+            mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
+
+            // Wait for reconnection attempts
+            try await Task.sleep(for: .seconds(4))
+
+            let endTime = Date()
+            let duration = endTime.timeIntervalSince(startTime)
+
+            // Verify we had enough time for reconnection attempts
+            #expect(duration >= 3.5)  // Should have time for 3 attempts with backoff
+
+            await transport.disconnect()
+        }
+
+        @Test("Heartbeat Failure Handling")
+        func testHeartbeatFailureHandling() async throws {
+            let mockConnection = MockNetworkConnection()
+
+            // Create transport with rapid heartbeats
+            let heartbeatConfig = NetworkTransport.HeartbeatConfiguration(
+                enabled: true,
+                interval: 0.1
+            )
+
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: heartbeatConfig
+            )
+
+            try await transport.connect()
+
+            // Wait for initial heartbeat
+            try await Task.sleep(for: .seconds(1.2))
+
+            // Simulate failure during heartbeat
+            mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
+
+            // Wait for potential recovery
+            try await Task.sleep(for: .milliseconds(500))
+
+            // Verify connection state
+            if case .failed = mockConnection.state {
+                // expected
+            } else {
+                Issue.record("Expected state to be failed")
+            }
+
+            await transport.disconnect()
+        }
+
+        @Test("Resource Cleanup")
+        func testResourceCleanup() async throws {
+            weak var weakConnection: MockNetworkConnection?
+
+            do {
+                let mockConnection = MockNetworkConnection()
+                weakConnection = mockConnection
+
+                // Create and use transport in a separate scope
+                let transport = NetworkTransport(
+                    mockConnection,
+                    heartbeatConfig: .disabled
+                )
+
+                try await transport.connect()
+                await transport.disconnect()
+            }
+
+            // Wait for potential async cleanup
+            try await Task.sleep(for: .milliseconds(100))
+
+            // Verify connection is cleaned up
+            #expect(weakConnection == nil, "Connection was not properly cleaned up")
+        }
+    }
+#endif

Tests/MCPTests/ServerTests.swift

@@ -39,7 +39,7 @@ struct ServerTests {
         try await server.start(transport: transport)
 
         // Wait for message processing and response
-        try await Task.sleep(nanoseconds: 100_000_000)  // 100ms
+        try await Task.sleep(for: .milliseconds(100))
 
         #expect(await transport.sentMessages.count == 1)
 
@@ -114,7 +114,7 @@ struct ServerTests {
         }
 
         // Wait for server to initialize
-        try await Task.sleep(nanoseconds: 10_000_000)  // 10ms
+        try await Task.sleep(for: .milliseconds(10))
 
         // Queue an initialize request from blocked client
         try await transport.queue(
@@ -127,7 +127,7 @@ struct ServerTests {
             ))
 
         // Wait for message processing
-        try await Task.sleep(nanoseconds: 200_000_000)  // 200ms
+        try await Task.sleep(for: .milliseconds(200))
 
         #expect(await transport.sentMessages.count >= 1)