Update watch stream client codes

client/client.go

@@ -431,16 +431,22 @@ func (c *Client) Watch(
 		return nil, err
 	}
 
-	pbResp := stream.Msg()
-	if err != nil {
-		return nil, err
+	for stream.Receive() {
+		pbResp := stream.Msg()
+		if err != nil {
+			return nil, err
+		}
+		if _, err := handleResponse(pbResp, doc); err != nil {
+			return nil, err
+		}
+		break
 	}
-	if _, err := handleResponse(pbResp, doc); err != nil {
-		return nil, err
+	if err = stream.Err(); err != nil {
+		return nil, connect.NewError(connect.CodeUnknown, err)
 	}
 
 	go func() {
-		for {
+		for stream.Receive() {
 			pbResp := stream.Msg()
 			if err != nil {
 				rch <- WatchResponse{Err: err}
@@ -459,6 +465,11 @@ func (c *Client) Watch(
 
 			rch <- *resp
 		}
+		if err = stream.Err(); err != nil {
+			rch <- WatchResponse{Err: err}
+			close(rch)
+			return
+		}
 	}()
 
 	// TODO(hackerwins): We need to revise the implementation of the watch

server/rpc/server_test.go

@@ -660,8 +660,11 @@ func TestSDKRPCServerBackend(t *testing.T) {
 		assert.NoError(t, err)
 
 		// check if stream is open
-		_ = watchResp.Msg()
-		assert.NoError(t, err)
+		for watchResp.Receive() {
+			resp := watchResp.Msg()
+			assert.NotNil(t, resp)
+			break
+		}
 
 		//// wait for MaxConnectionAge + MaxConnectionAgeGrace
 		//time.Sleep(helper.RPCMaxConnectionAge + helper.RPCMaxConnectionAgeGrace)