diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift index b4240e07..616e6c09 100644 --- a/Sources/MCP/Base/Transports/NetworkTransport.swift +++ b/Sources/MCP/Base/Transports/NetworkTransport.swift @@ -624,26 +624,27 @@ import Logging continue } - // Check if this is a heartbeat message - if Heartbeat.isHeartbeat(newData) { - logger.trace("Received heartbeat from peer") - - // Extract timestamp if available - if let heartbeat = Heartbeat.from(data: newData) { - logger.trace("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 - while let newlineIndex = buffer.firstIndex(of: UInt8(ascii: "\n")) { + // TCP reads may contain multiple frames or only part of a frame. + while !buffer.isEmpty { + if Heartbeat.isHeartbeat(buffer) { + // A heartbeat is four magic bytes followed by an eight-byte timestamp. + let heartbeatSize = 12 + guard buffer.count >= heartbeatSize else { break } + logger.trace("Received heartbeat from peer") + if let heartbeat = Heartbeat.from(data: Data(buffer.prefix(heartbeatSize))) { + logger.trace("Heartbeat timestamp: \(heartbeat.timestamp)") + } + buffer = buffer.dropFirst(heartbeatSize) + continue + } + + guard let newlineIndex = buffer.firstIndex(of: UInt8(ascii: "\n")) else { + break + } let messageData = buffer[.. 0 { + mockConnection.queueDataForReceiving(Data(payload.prefix(splitAt))) + } + mockConnection.queueDataForReceiving(Data(payload.dropFirst(splitAt))) + // A separate marker lets the unfixed implementation finish the test + // even when it discards the preceding messages. + mockConnection.queueDataForReceiving(marker + newline) + + try await transport.connect() + var received: [Data] = [] + do { + for try await data in await transport.receive() { + received.append(data) + if data == marker { break } + } + } catch { + await transport.disconnect() + throw error + } + await transport.disconnect() + + #expect(received == [first, second, marker]) + } + @Test("Heartbeat Send and Receive") func testNetworkTransportHeartbeat() async throws { let mockConnection = MockNetworkConnection()