diff --git a/ui.h b/ui.h
index bc728b0..0d5ab55 100644
--- a/ui.h
+++ b/ui.h
@@ -69,30 +69,27 @@
 
     // --- key handling ---
 
-    // Wait for keypress and return it.  May return -1 after timeout.
+    // Wait for a key and return it.  May return -1 after timeout.
     virtual int WaitKey();
 
     virtual bool IsKeyPressed(int key);
+    virtual bool IsLongPress();
 
     // Erase any queued-up keys.
     virtual void FlushKeys();
 
-    // Called on each keypress, even while operations are in progress.
+    // Called on each key press, even while operations are in progress.
     // Return value indicates whether an immediate operation should be
     // triggered (toggling the display, rebooting the device), or if
     // the key should be enqueued for use by the main thread.
     enum KeyAction { ENQUEUE, TOGGLE, REBOOT, IGNORE };
-    virtual KeyAction CheckKey(int key);
-
-    // Called immediately before each call to CheckKey(), tell you if
-    // the key was long-pressed.
-    virtual void NextCheckKeyIsLong(bool is_long_press);
+    virtual KeyAction CheckKey(int key, bool is_long_press);
 
     // Called when a key is held down long enough to have been a
     // long-press (but before the key is released).  This means that
     // if the key is eventually registered (released without any other
-    // keys being pressed in the meantime), NextCheckKeyIsLong() will
-    // be called with "true".
+    // keys being pressed in the meantime), CheckKey will be called with
+    // 'is_long_press' true.
     virtual void KeyLongPress(int key);
 
     // Normally in recovery there's a key sequence that triggers
@@ -110,8 +107,8 @@
     virtual void StartMenu(const char* const * headers, const char* const * items,
                            int initial_selection) = 0;
 
-    // Set the menu highlight to the given index, and return it (capped to
-    // the range [0..numitems).
+    // Set the menu highlight to the given index, wrapping if necessary.
+    // Returns the actual item selected.
     virtual int SelectMenu(int sel) = 0;
 
     // End menu mode, resetting the text overlay so that ui_print()
@@ -136,14 +133,20 @@
     int consecutive_power_keys;
     int last_key;
 
-    typedef struct {
+    bool has_power_key;
+    bool has_up_key;
+    bool has_down_key;
+
+    struct key_timer_t {
         RecoveryUI* ui;
         int key_code;
         int count;
-    } key_timer_t;
+    };
 
     pthread_t input_t;
 
+    void OnKeyDetected(int key_code);
+
     static void* input_thread(void* cookie);
     static int input_callback(int fd, uint32_t epevents, void* data);
     void process_key(int key_code, int updown);