From 051030c96dd50a0e5ff7ae9e5f2913321c40fdc1 Mon Sep 17 00:00:00 2001
From: Vishal Punjabi <10079537+vshalpnjabi@users.noreply.github.com>
Date: Wed, 28 Oct 2026 18:18:03 -0700
Subject: [PATCH] mcp: cancel pending keepalive ping when the session closes

startKeepalive derived each ping's context from context.Background,
so cancelling the keepalive did not reach a ping already in flight.
Closing a session whose peer had stopped answering, such as on a
client DELETE, then waited up to KeepAlive/2 for the ping to time out
and logged it as a keepalive failure.

Derive the ping context from the keepalive context, and stop the loop
without counting a failure when that context has been cancelled.

Fixes #1470
---
 mcp/mcp_test.go | 56 +++++++++++++++++++++++++++++++++++++++++++++++++
 mcp/shared.go   |  5 ++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go
index b15cd7fb..a9b825bb 100644
--- a/mcp/mcp_test.go
+++ b/mcp/mcp_test.go
@@ -2180,6 +2180,62 @@ func TestStartKeepalive_SuccessResetsFailures(t *testing.T) {
 	})
 }
 
+// blockingKeepaliveSession is a keepaliveSession test double whose Ping
+// blocks until its context is done, like a ping the peer never answers.
+type blockingKeepaliveSession struct {
+	pingsDone  atomic.Int64
+	closeCalls atomic.Int64
+}
+
+func (s *blockingKeepaliveSession) Ping(ctx context.Context, _ *PingParams) error {
+	<-ctx.Done()
+	s.pingsDone.Add(1)
+	return ctx.Err()
+}
+
+func (s *blockingKeepaliveSession) Close() error {
+	s.closeCalls.Add(1)
+	return nil
+}
+
+// TestStartKeepalive_CancelStopsPendingPing verifies that cancelling the
+// keepalive, as closing the session does, ends a pending ping at once rather
+// than after its timeout, and is not reported as a ping failure.
+// Regression test for #1470.
+func TestStartKeepalive_CancelStopsPendingPing(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		const interval = 100 * time.Millisecond
+		sess := &blockingKeepaliveSession{}
+		var buf bytes.Buffer
+		logger := slog.New(slog.NewTextHandler(&buf, nil))
+		var cancel context.CancelFunc
+		startKeepalive(sess, interval, 1, &cancel, logger)
+
+		// Let the first ping start; the peer never answers it.
+		time.Sleep(interval + interval/4)
+		synctest.Wait()
+		if got := sess.pingsDone.Load(); got != 0 {
+			t.Fatalf("ping returned before cancel: pingsDone=%d", got)
+		}
+
+		cancel()
+		synctest.Wait()
+		if got := sess.pingsDone.Load(); got != 1 {
+			t.Errorf("pending ping not stopped by cancel: pingsDone=%d", got)
+		}
+
+		// Run past the ping timeout and the next tick: nothing more happens.
+		time.Sleep(2 * interval)
+		synctest.Wait()
+		if got := sess.closeCalls.Load(); got != 0 {
+			t.Errorf("session closed after cancel: closeCalls=%d", got)
+		}
+		if got := buf.String(); got != "" {
+			t.Errorf("unexpected log output after cancel:\n%s", got)
+		}
+	})
+}
+
 func TestAddTool_DuplicateNoPanicAndNoDuplicate(t *testing.T) {
 	// Adding the same tool pointer twice should not panic and should not
 	// produce duplicates in the server's tool list.
diff --git a/mcp/shared.go b/mcp/shared.go
index 7f5a6abe..23ccc85c 100644
--- a/mcp/shared.go
+++ b/mcp/shared.go
@@ -991,9 +991,12 @@ func startKeepalive(session keepaliveSession, interval time.Duration, failureThr
 			case <-ctx.Done():
 				return
 			case <-ticker.C:
-				pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2)
+				pingCtx, pingCancel := context.WithTimeout(ctx, interval/2)
 				err := session.Ping(pingCtx, nil)
 				pingCancel()
+				if ctx.Err() != nil {
+					return
+				}
 				if err == nil {
 					consecutiveFailures = 0
 					continue