add missing include after refactoring