Improve connection logic