Fix wait_on_exit to wait for keys in popup and not only in main document