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
This commit is contained in:
Mattt
GitHub
parent
266c02d2d1
commit
87f33d022a
3 changed files with 1224 additions and 32 deletions
|
|
@ -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
|
||||
/// Logger instance for transport-related events
|
||||
public nonisolated let logger: Logger
|
||||
/// Represents a heartbeat message for connection health monitoring.
|
||||
public struct Heartbeat: RawRepresentable, Hashable, Sendable {
|
||||
/// 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:
|
||||
// Wait for ready or failed state
|
||||
break
|
||||
case .waiting(let error):
|
||||
self.logger.debug("Connection waiting: \(error)")
|
||||
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
|
||||
if connection.state != .ready {
|
||||
connection.start(queue: .main)
|
||||
} else {
|
||||
Task { @MainActor in
|
||||
await self.handleConnectionReady(continuation: continuation)
|
||||
}
|
||||
}
|
||||
connection.start(queue: .main)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -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).")
|
||||
break // Exit loop gracefully
|
||||
consecutiveEmptyReads += 1
|
||||
|
||||
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
|
||||
|
|
|
|||
688
Tests/MCPTests/NetworkTransportTests.swift
Normal file
688
Tests/MCPTests/NetworkTransportTests.swift
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue