From 5bfaeca058f57aa637b23195e58c44efbd015e95 Mon Sep 17 00:00:00 2001
From: bbdu3 <ergouyang854@gamil.com>
Date: Fri, 6 Nov 2026 14:06:28 +0800
Subject: [PATCH 1/2] Fix JSON-RPC message loss when TCP reads contain
 heartbeats

Signed-off-by: bbdu3 <ergouyang854@gamil.com>
---
 .../Base/Transports/NetworkTransport.swift    | 33 +++++++-------
 Tests/MCPTests/NetworkTransportTests.swift    | 43 +++++++++++++++++++
 2 files changed, 60 insertions(+), 16 deletions(-)

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[..<newlineIndex]
                         buffer = buffer[(newlineIndex + 1)...]
 
diff --git a/Tests/MCPTests/NetworkTransportTests.swift b/Tests/MCPTests/NetworkTransportTests.swift
index 9c15577c..13cbf56f 100644
--- a/Tests/MCPTests/NetworkTransportTests.swift
+++ b/Tests/MCPTests/NetworkTransportTests.swift
@@ -341,6 +341,49 @@ import Testing
             await transport.disconnect()
         }
 
+        @Test("Receive messages coalesced with heartbeats", .timeLimit(.minutes(1)), arguments: 0...13)
+        func testReceiveMessagesCoalescedWithHeartbeats(splitAt: Int) async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .disabled,
+                reconnectionConfig: .disabled
+            )
+            let heartbeat = NetworkTransport.Heartbeat(
+                timestamp: Date(timeIntervalSinceReferenceDate: 10)
+            ).data
+            let first = Data(#"{"jsonrpc":"2.0","id":1,"result":{}}"#.utf8)
+            let second = Data(#"{"jsonrpc":"2.0","id":2,"result":{}}"#.utf8)
+            let marker = Data(#"{"jsonrpc":"2.0","id":3,"result":{}}"#.utf8)
+            let newline = Data([0x0A])
+            let payload = heartbeat + first + newline + heartbeat + second + newline
+
+            // Zero delivers one coalesced chunk; the other cases split the first
+            // heartbeat's header/timestamp, or the start of the following JSON.
+            if splitAt > 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()

From 63587287a66c51eb2fea2abcbcb4650022db53f7 Mon Sep 17 00:00:00 2001
From: bbdu3 <ergouyang854@gamil.com>
Date: Mon, 9 Nov 2026 10:36:14 +0800
Subject: [PATCH 2/2] fix ci runner

Signed-off-by: bbdu3 <ergouyang854@gamil.com>
---
 .github/workflows/ci.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ba30a965..2aa8e355 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -25,7 +25,7 @@ jobs:
           # Current leg: catches breakage from new Xcode and SDK releases.
           - os: macos-latest
             swift-version: Xcode default
-          - os: ubuntu-latest
+          - os: ubuntu-24.04
             swift-version: 6.1.0
 
     runs-on: ${{ matrix.os }}
@@ -35,13 +35,13 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Prepare Linux apt dependencies
-        if: matrix.os == 'ubuntu-latest'
+        if: matrix.os == 'ubuntu-24.04'
         run: |
           sudo -n apt-get update
           sudo -n apt-get install -y --no-install-recommends libcurl4-openssl-dev
 
       - name: Setup Swift on Linux
-        if: matrix.os == 'ubuntu-latest'
+        if: matrix.os == 'ubuntu-24.04'
         uses: vapor/swiftly-action@v0.2
         with:
           toolchain: ${{ matrix.swift-version }}