Error handling improvement