Add additional error handling