Fix data races in NetworkTransport and MockNetworkConnection

Sources/MCP/Base/Transports/NetworkTransport.swift

@@ -348,12 +348,7 @@ import Logging
         ///
         /// - Throws: Error if the connection fails
         private func attemptConnection() async throws {
-            // Create stream for state changes with proper cleanup
             let stateStream = AsyncStream<NWConnection.State> { continuation in
-                continuation.onTermination = { [weak self] _ in
-                    self?.connection.stateUpdateHandler = nil
-                }
-
                 connection.stateUpdateHandler = { state in
                     continuation.yield(state)
                     switch state {
@@ -367,6 +362,10 @@ import Logging
 
             connection.start(queue: .main)
 
+            // Clear the handler on the actor's executor once the stream is exhausted,
+            // regardless of whether the function returns normally or throws.
+            defer { connection.stateUpdateHandler = nil }
+
             // Process states until terminal state
             for await state in stateStream {
                 switch state {

Tests/MCPTests/NetworkTransportTests.swift

@@ -9,6 +9,8 @@ import Testing
 
     /// A mock implementation of NetworkConnectionProtocol for testing
     final class MockNetworkConnection: NetworkConnectionProtocol, @unchecked Sendable {
+        private let lock = NSLock()
+
         /// Current state of the connection
         private var mockState: NWConnection.State = .setup
 
@@ -20,26 +22,25 @@ import Testing
         private var sentData: [Data] = []
 
         /// The state update handler
-        public var stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))?
+        public var stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))? {
+            get { lock.withLock { _stateUpdateHandler } }
+            set { lock.withLock { _stateUpdateHandler = newValue } }
+        }
+        private var _stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))?
 
         /// Current state
         var state: NWConnection.State {
-            return mockState
+            lock.withLock { mockState }
         }
 
         /// Initialize a mock connection
         init() {}
 
         /// Start the connection
         func start(queue: DispatchQueue) {
-            // Simulate successful connection by default
-            Task { @MainActor in
-                if mockError != nil {
-                    self.updateState(mockState)
-                } else {
-                    self.updateState(.ready)
-                }
-            }
+            let (error, state) = lock.withLock { (mockError, mockState) }
+            let nextState: NWConnection.State = error != nil ? state : .ready
+            updateState(nextState)
         }
 
         /// Send data through the connection
@@ -49,15 +50,14 @@ import Testing
             isComplete: Bool,
             completion: NWConnection.SendCompletion
         ) {
-            if let content = content {
-                sentData.append(content)
+            let error = lock.withLock {
+                if let content = content { sentData.append(content) }
+                return mockError
             }
 
             switch completion {
             case .contentProcessed(let handler):
-                Task { @MainActor in
-                    handler(self.mockError as? NWError)
-                }
+                handler(error as? NWError)
             default:
                 break
             }
@@ -71,28 +71,21 @@ import Testing
                 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 (state, error, data): (NWConnection.State, Swift.Error?, Data?) = lock.withLock {
+                let d = dataToReceive.isEmpty ? nil : dataToReceive.removeFirst()
+                return (mockState, mockError, d)
+            }
 
-                let data = self.dataToReceive.removeFirst()
-                completion(data, nil, self.dataToReceive.isEmpty, nil)
+            if state == .cancelled {
+                completion(nil, nil, true, NWError.posix(POSIXErrorCode.ECANCELED))
+                return
+            }
+            if let error {
+                completion(nil, nil, false, error as? NWError)
+                return
             }
+            let isLast = lock.withLock { dataToReceive.isEmpty }
+            completion(data ?? Data(), nil, data == nil ? false : isLast, nil)
         }
 
         /// Cancel the connection
@@ -113,10 +106,8 @@ import Testing
         }
 
         /// Simulate a connection failure
-        func simulateFailure(
-            error: Swift.Error? = nil
-        ) {
-            mockError = error
+        func simulateFailure(error: Swift.Error? = nil) {
+            lock.withLock { mockError = error }
             if let nwError = error as? NWError {
                 updateState(.failed(nwError))
             } else {
@@ -131,15 +122,17 @@ import Testing
 
         /// Update the connection state and notify handler
         private func updateState(_ newState: NWConnection.State) {
-            mockState = newState
-            Task { @MainActor in
-                self.stateUpdateHandler?(newState)
+            // Snapshot handler outside the lock to avoid holding the lock during callback
+            let handler = lock.withLock { () -> (@Sendable (NWConnection.State) -> Void)? in
+                mockState = newState
+                return _stateUpdateHandler
             }
+            handler?(newState)
         }
 
         /// Queue data to be received
         func queueDataForReceiving(_ data: Data) {
-            dataToReceive.append(data)
+            lock.withLock { dataToReceive.append(data) }
         }
 
         /// Queue a heartbeat message to be received
@@ -162,12 +155,12 @@ import Testing
 
         /// Get all sent data
         func getSentData() -> [Data] {
-            return sentData
+            lock.withLock { sentData }
         }
 
         /// Clear sent data buffer
         func clearSentData() {
-            sentData.removeAll()
+            lock.withLock { sentData.removeAll() }
         }
     }