update

Author: dugenkui03

client/http_test.go

@@ -13,10 +13,10 @@ func TestHTTPClient(t *testing.T) {
 	hooks := &server.Hooks{}
 	hooks.AddAfterCallTool(func(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult) {
 		clientSession := server.ClientSessionFromContext(ctx)
-		// return until all the notifications are handled
+		// wait until all the notifications are handled
 		for len(clientSession.NotificationChannel()) > 0 {
 		}
-		time.Sleep(time.Millisecond * 200)
+		time.Sleep(time.Millisecond * 50)
 	})
 
 	// Create MCP server with capabilities
@@ -58,11 +58,10 @@ func TestHTTPClient(t *testing.T) {
 		},
 	)
 
-	// Initialize
 	testServer := server.NewTestStreamableHTTPServer(mcpServer)
 	defer testServer.Close()
 
-	t.Run("Can create client", func(t *testing.T) {
+	t.Run("Can receive notification from server", func(t *testing.T) {
 		client, err := NewStreamableHttpClient(testServer.URL)
 		if err != nil {
 			t.Fatalf("create client failed %v", err)

server/streamable_http.go

@@ -246,7 +246,6 @@ func (s *StreamableHTTPServer) handlePost(w http.ResponseWriter, r *http.Request
 	mu := sync.Mutex{}
 	upgradedHeader := false
 	done := make(chan struct{})
-	defer close(done)
 
 	go func() {
 		for {
@@ -255,6 +254,12 @@ func (s *StreamableHTTPServer) handlePost(w http.ResponseWriter, r *http.Request
 				func() {
 					mu.Lock()
 					defer mu.Unlock()
+					// if the done chan is closed, as the request is terminated, just return
+					select {
+					case <-done:
+						return
+					default:
+					}
 					defer func() {
 						flusher, ok := w.(http.Flusher)
 						if ok {
@@ -295,6 +300,8 @@ func (s *StreamableHTTPServer) handlePost(w http.ResponseWriter, r *http.Request
 	// Write response
 	mu.Lock()
 	defer mu.Unlock()
+	// close the done chan before unlock
+	defer close(done)
 	if ctx.Err() != nil {
 		return
 	}