Fix bugs in wait