From bef2f0716bb50e0099314449d51498131a3b932e Mon Sep 17 00:00:00 2001 From: Alex Bullitt Date: Wed, 25 Mar 2026 19:19:54 -0700 Subject: [PATCH 1/2] fix: eliminate data races in NetworkTransport and MockNetworkConnection NetworkTransport.attemptConnection() was writing connection.stateUpdateHandler from two contexts: the actor's executor (setting the handler) and AsyncStream's onTermination callback (clearing it to nil), which can fire on any thread. Fix: remove onTermination and use defer to clear the handler on the actor's executor after the for-await loop exits, regardless of normal return or throw. MockNetworkConnection was @unchecked Sendable with entirely unprotected mutable state. The Task { @MainActor in } wrappers were inconsistent: simulateFailure wrote mockError directly while receive read it on @MainActor, and stateUpdateHandler was set from NetworkTransport's actor executor while being read/written on @MainActor. Fix: add NSLock protecting all mutable state; snapshot handler before calling it to avoid holding the lock during the callback. --- .../Base/Transports/NetworkTransport.swift | 9 +- Tests/MCPTests/NetworkTransportTests.swift | 83 +++++++++---------- 2 files changed, 42 insertions(+), 50 deletions(-) diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift index 7c4d00f3..68f9856e 100644 --- a/Sources/MCP/Base/Transports/NetworkTransport.swift +++ b/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 { 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 { diff --git a/Tests/MCPTests/NetworkTransportTests.swift b/Tests/MCPTests/NetworkTransportTests.swift index 9c02c90c..b5a18b41 100644 --- a/Tests/MCPTests/NetworkTransportTests.swift +++ b/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,11 +22,15 @@ 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 @@ -32,14 +38,9 @@ import Testing /// 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() } } } From 419a7b61765c276bc02af0387827566069bcaa19 Mon Sep 17 00:00:00 2001 From: Varun Date: Sat, 31 Oct 2026 06:31:39 -0500 Subject: [PATCH 2/2] test: dequeue mock receive data under one lock, after state checks MockNetworkConnection.receive removed a queued item before checking for a cancelled or failed connection, dropping that item on those paths, and read isLast under a separate lock acquisition. Check state, dequeue, and compute isLast together under a single lock. Addresses review feedback on #301. --- Tests/MCPTests/NetworkTransportTests.swift | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/Tests/MCPTests/NetworkTransportTests.swift b/Tests/MCPTests/NetworkTransportTests.swift index b5a18b41..297356d4 100644 --- a/Tests/MCPTests/NetworkTransportTests.swift +++ b/Tests/MCPTests/NetworkTransportTests.swift @@ -71,10 +71,14 @@ import Testing Data?, NWConnection.ContentContext?, Bool, NWError? ) -> Void ) { - let (state, error, data): (NWConnection.State, Swift.Error?, Data?) = lock.withLock { - let d = dataToReceive.isEmpty ? nil : dataToReceive.removeFirst() - return (mockState, mockError, d) - } + let (state, error, data, isLast): (NWConnection.State, Swift.Error?, Data?, Bool) = + lock.withLock { + guard mockState != .cancelled, mockError == nil, !dataToReceive.isEmpty else { + return (mockState, mockError, nil, false) + } + let data = dataToReceive.removeFirst() + return (mockState, mockError, data, dataToReceive.isEmpty) + } if state == .cancelled { completion(nil, nil, true, NWError.posix(POSIXErrorCode.ECANCELED)) @@ -84,8 +88,7 @@ import Testing completion(nil, nil, false, error as? NWError) return } - let isLast = lock.withLock { dataToReceive.isEmpty } - completion(data ?? Data(), nil, data == nil ? false : isLast, nil) + completion(data ?? Data(), nil, isLast, nil) } /// Cancel the connection