diff --git a/recovery_ui/stub_ui.cpp b/recovery_ui/stub_ui.cpp
index a56b3f7..87605cf 100644
--- a/recovery_ui/stub_ui.cpp
+++ b/recovery_ui/stub_ui.cpp
@@ -25,11 +25,12 @@
                                 size_t /* initial_selection */, bool /* menu_only */,
                                 const std::function<int(int, bool)>& /*key_handler*/) {
   while (true) {
-    int key = WaitKey();
     // Exit the loop in the case of interruption or time out.
-    if (key == static_cast<int>(KeyError::INTERRUPTED) ||
-        key == static_cast<int>(KeyError::TIMED_OUT)) {
-      return static_cast<size_t>(key);
+    InputEvent evt = WaitInputEvent();
+    if (evt.type() == EventType::EXTRA) {
+      if (evt.key() == static_cast<int>(KeyError::INTERRUPTED) ||
+        evt.key() == static_cast<int>(KeyError::TIMED_OUT))
+          return static_cast<size_t>(evt.key());
     }
   }
   LOG(FATAL) << "Unreachable key selected in ShowMenu of stub UI";