Fixing the wait for endpoint logic