just added some new feature. Might have a bug too. Can u review it?

I added a new awesome feature but i think i caused a bug, can you see if you can find and fix the bug while keeping my feature addition intact? :D