Fix connection handling