From c0ff1e6ff925ec3928121c4d5f60a4c4180e8595 Mon Sep 17 00:00:00 2001
From: EhsanAzish80 <ehzish@3nsofts.com>
Date: Sun, 8 Nov 2026 12:00:40 +0300
Subject: [PATCH] fix: stop heartbeats after failed connection closes

---
 .../Base/Transports/NetworkTransport.swift    |  3 ++
 Tests/MCPTests/NetworkTransportTests.swift    | 28 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift
index b4240e07..fbcde804 100644
--- a/Sources/MCP/Base/Transports/NetworkTransport.swift
+++ b/Sources/MCP/Base/Transports/NetworkTransport.swift
@@ -612,6 +612,9 @@ import Logging
                                 if !isStopping {
                                     // Release the underlying socket and finish the stream so
                                     // receive() callers don't hang.
+                                    isConnected = false
+                                    heartbeatTask?.cancel()
+                                    heartbeatTask = nil
                                     connection.cancel()
                                     messageContinuation.finish(throwing: MCPError.connectionClosed)
                                 }
diff --git a/Tests/MCPTests/NetworkTransportTests.swift b/Tests/MCPTests/NetworkTransportTests.swift
index 9c15577c..9acf85d4 100644
--- a/Tests/MCPTests/NetworkTransportTests.swift
+++ b/Tests/MCPTests/NetworkTransportTests.swift
@@ -668,6 +668,34 @@ import Testing
             await transport.disconnect()
         }
 
+        @Test("Heartbeats stop when empty reads reveal a failed connection")
+        func testHeartbeatsStopAfterEmptyReadFailure() async throws {
+            let mockConnection = MockNetworkConnection()
+            let transport = NetworkTransport(
+                mockConnection,
+                heartbeatConfig: .init(enabled: true, interval: 0.1)
+            )
+
+            try await transport.connect()
+            mockConnection.simulateFailure()
+
+            let deadline = ContinuousClock.now + .seconds(3)
+            while mockConnection.state != .cancelled && ContinuousClock.now < deadline {
+                try await Task.sleep(for: .milliseconds(50))
+            }
+            #expect(mockConnection.state == .cancelled)
+
+            await #expect(throws: MCPError.self) {
+                try await transport.send(Data())
+            }
+
+            let heartbeatsSent = mockConnection.getSentData().count
+            try await Task.sleep(for: .seconds(2.5))
+            #expect(mockConnection.getSentData().count == heartbeatsSent)
+
+            await transport.disconnect()
+        }
+
         @Test("Resource Cleanup")
         func testResourceCleanup() async throws {
             weak var weakConnection: MockNetworkConnection?