Fixing support for nested while loops with callbacks + tests