chore: address review comments

Author: rwjblue-glean

mcp/client.go

@@ -148,6 +148,11 @@ type ClientSession struct {
 // from being handled, and waiting for ongoing requests to return. Close then
 // terminates the connection.
 func (cs *ClientSession) Close() error {
+	// Note: keepaliveCancel access is safe without a mutex because:
+	// 1. keepaliveCancel is only written once during startKeepalive (happens-before user Close calls)
+	// 2. context.CancelFunc is safe to call multiple times and from multiple goroutines
+	// 3. The keepalive goroutine calls Close on ping failure, but this is safe since
+	//    Close is idempotent and conn.Close() handles concurrent calls correctly
 	if cs.keepaliveCancel != nil {
 		cs.keepaliveCancel()
 	}
@@ -162,29 +167,7 @@ func (cs *ClientSession) Wait() error {
 
 // startKeepalive starts the keepalive mechanism for this client session.
 func (cs *ClientSession) startKeepalive(interval time.Duration) {
-	ctx, cancel := context.WithCancel(context.Background())
-	cs.keepaliveCancel = cancel
-
-	go func() {
-		ticker := time.NewTicker(interval)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-ticker.C:
-				pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2)
-				err := cs.Ping(pingCtx, nil)
-				pingCancel()
-				if err != nil {
-					// Ping failed, close the session
-					_ = cs.Close()
-					return
-				}
-			}
-		}
-	}()
+	startKeepalive(cs, interval, &cs.keepaliveCancel)
 }
 
 // AddRoots adds the given roots to the client,

mcp/mcp_test.go

@@ -840,6 +840,7 @@ func nopHandler(context.Context, *ServerSession, *CallToolParamsFor[map[string]a
 }
 
 func TestKeepAlive(t *testing.T) {
+	// TODO: try to use the new synctest package for this test.
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 

mcp/server.go

@@ -690,6 +690,11 @@ func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (*e
 // Close then terminates the connection.
 func (ss *ServerSession) Close() error {
 	if ss.keepaliveCancel != nil {
+		// Note: keepaliveCancel access is safe without a mutex because:
+		// 1. keepaliveCancel is only written once during startKeepalive (happens-before user Close calls)
+		// 2. context.CancelFunc is safe to call multiple times and from multiple goroutines
+		// 3. The keepalive goroutine calls Close on ping failure, but this is safe since
+		//    Close is idempotent and conn.Close() handles concurrent calls correctly
 		ss.keepaliveCancel()
 	}
 	return ss.conn.Close()
@@ -702,29 +707,7 @@ func (ss *ServerSession) Wait() error {
 
 // startKeepalive starts the keepalive mechanism for this server session.
 func (ss *ServerSession) startKeepalive(interval time.Duration) {
-	ctx, cancel := context.WithCancel(context.Background())
-	ss.keepaliveCancel = cancel
-
-	go func() {
-		ticker := time.NewTicker(interval)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-ticker.C:
-				pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2)
-				err := ss.Ping(pingCtx, nil)
-				pingCancel()
-				if err != nil {
-					// Ping failed, close the session
-					_ = ss.Close()
-					return
-				}
-			}
-		}
-	}()
+	startKeepalive(ss, interval, &ss.keepaliveCancel)
 }
 
 // pageToken is the internal structure for the opaque pagination cursor.

mcp/shared.go

@@ -308,3 +308,38 @@ type listResult[T any] interface {
 	// Returns a pointer to the param's NextCursor field.
 	nextCursorPtr() *string
 }
+
+// keepaliveSession represents a session that supports keepalive functionality.
+type keepaliveSession interface {
+	Ping(ctx context.Context, params *PingParams) error
+	Close() error
+}
+
+// startKeepalive starts the keepalive mechanism for a session.
+// It assigns the cancel function to the provided cancelPtr and starts a goroutine
+// that sends ping messages at the specified interval.
+func startKeepalive(session keepaliveSession, interval time.Duration, cancelPtr *context.CancelFunc) {
+	ctx, cancel := context.WithCancel(context.Background())
+	*cancelPtr = cancel
+
+	go func() {
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2)
+				err := session.Ping(pingCtx, nil)
+				pingCancel()
+				if err != nil {
+					// Ping failed, close the session
+					_ = session.Close()
+					return
+				}
+			}
+		}
+	}()
+}