1 files changed, 14 insertions, 3 deletions
diff --git a/src/Kaleidoscope-LayerHighlighter.cpp b/src/Kaleidoscope-LayerHighlighter.cpp
index 3ed56cb..7c30761 100644
--- a/src/Kaleidoscope-LayerHighlighter.cpp
+++ b/src/Kaleidoscope-LayerHighlighter.cpp
@@ -8,14 +8,25 @@ kaleidoscope::EventHandlerResult LayerHighlighter::onSetup(void) {
   return kaleidoscope::EventHandlerResult::OK;
 }
 
+kaleidoscope::EventHandlerResult LayerHighlighter::onKeyswitchEvent(Key &mappedKey, byte row, byte col, uint8_t keyState) {
+  if (mappedKey != LockLayer(layer)) {
+    return kaleidoscope::EventHandlerResult::OK;
+  }
+
+  if (keyToggledOn(keyState)) {
+    savedLEDMode = LEDControl.get_mode_index();
+  } else if (keyToggledOff(keyState)) {
+    LEDControl.set_mode(LEDControl.get_mode_index());
+  }
+
+  return kaleidoscope::EventHandlerResult::OK;
+}
+
 kaleidoscope::EventHandlerResult LayerHighlighter::afterEachCycle() {
   if (!Layer.isOn(layer)) {
-    LEDControl.set_mode(LEDControl.get_mode_index());
     return kaleidoscope::EventHandlerResult::OK;
   }
 
-  LEDControl.set_mode(LEDControl.get_mode_index());
-
   for (uint8_t r = 0; r < ROWS; r++) {
     for (uint8_t c = 0; c < COLS; c++) {
       Key k = Layer.lookupOnActiveLayer(r, c);