Added notification handling