fix data race in test

Author: dugenkui03

client/sse_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/stretchr/testify/assert"
 	"net/http"
+	"sync"
 	"testing"
 	"time"
 
@@ -460,16 +461,19 @@ func TestSSEMCPClient(t *testing.T) {
 			t.Fatalf("Failed to create client: %v", err)
 		}
 
+		mu := sync.Mutex{}
 		notificationNum := 0
 		var messageNotification *mcp.JSONRPCNotification
 		progressNotifications := make([]*mcp.JSONRPCNotification, 0)
 		client.OnNotification(func(notification mcp.JSONRPCNotification) {
+			mu.Lock()
 			if notification.Method == string(mcp.MethodNotificationMessage) {
 				messageNotification = &notification
 			} else if notification.Method == string(mcp.MethodNotificationProgress) {
 				progressNotifications = append(progressNotifications, &notification)
 			}
 			notificationNum += 1
+			mu.Unlock()
 		})
 		defer client.Close()
 
@@ -516,6 +520,7 @@ func TestSSEMCPClient(t *testing.T) {
 
 		time.Sleep(time.Millisecond * 200)
 
+		mu.Lock()
 		assert.Equal(t, notificationNum, 3)
 		assert.NotNil(t, messageNotification)
 		assert.Equal(t, messageNotification.Method, string(mcp.MethodNotificationMessage))
@@ -537,6 +542,7 @@ func TestSSEMCPClient(t *testing.T) {
 		assert.EqualValues(t, 100, progressNotifications[1].Params.AdditionalFields["progress"])
 		assert.Equal(t, "progress_token", progressNotifications[1].Params.AdditionalFields["progressToken"])
 		assert.EqualValues(t, 100, progressNotifications[1].Params.AdditionalFields["total"])
+		mu.Lock()
 	})
 
 	t.Run("Ensure the server does not send notifications", func(t *testing.T) {
@@ -545,9 +551,12 @@ func TestSSEMCPClient(t *testing.T) {
 			t.Fatalf("Failed to create client: %v", err)
 		}
 
+		mu := sync.Mutex{}
 		notifications := make([]*mcp.JSONRPCNotification, 0)
 		client.OnNotification(func(notification mcp.JSONRPCNotification) {
+			mu.Lock()
 			notifications = append(notifications, &notification)
+			mu.Unlock()
 		})
 		defer client.Close()
 
@@ -585,7 +594,9 @@ func TestSSEMCPClient(t *testing.T) {
 		_, _ = client.CallTool(ctx, request)
 		time.Sleep(time.Millisecond * 200)
 
+		mu.Lock()
 		assert.Len(t, notifications, 0)
+		mu.Unlock()
 	})
 
 	t.Run("GetPrompt for testing log and progress notification", func(t *testing.T) {
@@ -594,17 +605,19 @@ func TestSSEMCPClient(t *testing.T) {
 			t.Fatalf("Failed to create client: %v", err)
 		}
 
+		mu := sync.Mutex{}
 		var messageNotification *mcp.JSONRPCNotification
 		progressNotifications := make([]*mcp.JSONRPCNotification, 0)
 		notificationNum := 0
 		client.OnNotification(func(notification mcp.JSONRPCNotification) {
-			println(notification.Method)
+			mu.Lock()
 			if notification.Method == string(mcp.MethodNotificationMessage) {
 				messageNotification = &notification
 			} else if notification.Method == string(mcp.MethodNotificationProgress) {
 				progressNotifications = append(progressNotifications, &notification)
 			}
 			notificationNum += 1
+			mu.Unlock()
 		})
 		defer client.Close()
 
@@ -645,6 +658,7 @@ func TestSSEMCPClient(t *testing.T) {
 		if err != nil {
 			t.Fatalf("GetPrompt failed: %v", err)
 		}
+		mu.Lock()
 		assert.NotNil(t, result)
 		assert.Len(t, result.Messages, 1)
 		assert.Equal(t, result.Messages[0].Role, mcp.RoleAssistant)
@@ -675,6 +689,7 @@ func TestSSEMCPClient(t *testing.T) {
 		assert.EqualValues(t, 100, progressNotifications[1].Params.AdditionalFields["progress"])
 		assert.Equal(t, "progress_token", progressNotifications[1].Params.AdditionalFields["progressToken"])
 		assert.EqualValues(t, 100, progressNotifications[1].Params.AdditionalFields["total"])
+		mu.Unlock()
 	})
 
 	t.Run("GetResource for testing log and progress notification", func(t *testing.T) {
@@ -683,17 +698,19 @@ func TestSSEMCPClient(t *testing.T) {
 			t.Fatalf("Failed to create client: %v", err)
 		}
 
+		mu := sync.Mutex{}
 		var messageNotification *mcp.JSONRPCNotification
 		progressNotifications := make([]*mcp.JSONRPCNotification, 0)
 		notificationNum := 0
 		client.OnNotification(func(notification mcp.JSONRPCNotification) {
-			println(notification.Method)
+			mu.Lock()
 			if notification.Method == string(mcp.MethodNotificationMessage) {
 				messageNotification = &notification
 			} else if notification.Method == string(mcp.MethodNotificationProgress) {
 				progressNotifications = append(progressNotifications, &notification)
 			}
 			notificationNum += 1
+			mu.Unlock()
 		})
 		defer client.Close()
 
@@ -735,6 +752,7 @@ func TestSSEMCPClient(t *testing.T) {
 			t.Fatalf("ReadResource failed: %v", err)
 		}
 
+		mu.Lock()
 		assert.NotNil(t, result)
 		assert.Len(t, result.Contents, 1)
 		assert.Equal(t, result.Contents[0].(mcp.TextResourceContents).URI, "resource://testresource")
@@ -763,5 +781,6 @@ func TestSSEMCPClient(t *testing.T) {
 		assert.EqualValues(t, 100, progressNotifications[1].Params.AdditionalFields["progress"])
 		assert.Equal(t, "progress_token", progressNotifications[1].Params.AdditionalFields["progressToken"])
 		assert.EqualValues(t, 100, progressNotifications[1].Params.AdditionalFields["total"])
+		mu.Unlock()
 	})
 }

server/sse.go

@@ -519,7 +519,8 @@ func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
 			var message string
 			if eventData, err := json.Marshal(response); err != nil {
 				// If there is an error marshalling the response, send a generic error response
-				log.Printf("failed to marshal response: %v", err)
+				marshal, _ := json.Marshal(response)
+				log.Printf("failed to marshal response: %v, response %s", err, string(marshal))
 				message = "event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n"
 			} else {
 				message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData)