ui: preserve custom button presses

openpilot/selfdrive/ui/layouts/main.py

@@ -28,6 +28,7 @@ class MainLayout(Widget):
     super().__init__()
 
     self._pm = messaging.PubMaster(['bookmarkButton', 'userBookmark'])
+    self._custom_button_sock = messaging.sub_sock('carState')
 
     self._sidebar = Sidebar()
     self._current_mode = MainState.HOME
@@ -60,7 +61,7 @@ class MainLayout(Widget):
       gui_app.push_widget(self._onboarding_window)
 
   def _render(self, _):
-    handle_custom_button(ui_state.sm, ui_state.params, self._custom_button_callbacks)
+    handle_custom_button(messaging.drain_sock(self._custom_button_sock), ui_state.params, self._custom_button_callbacks)
     self._handle_onroad_transition()
     self._render_main_content()
 

openpilot/selfdrive/ui/mici/layouts/main.py

@@ -24,6 +24,7 @@ class MiciMainLayout(Scroller):
     super().__init__(snap_items=True, spacing=0, pad=0, scroll_indicator=False, edge_shadows=False)
 
     self._pm = messaging.PubMaster(['bookmarkButton', 'userBookmark'])
+    self._custom_button_sock = messaging.sub_sock('carState')
 
     self._prev_onroad = False
     self._prev_standstill = False
@@ -100,7 +101,7 @@ class MiciMainLayout(Scroller):
     self._alerts_layout._update_state()
 
   def _render(self, _):
-    handle_custom_button(ui_state.sm, ui_state.params, self._custom_button_callbacks)
+    handle_custom_button(messaging.drain_sock(self._custom_button_sock), ui_state.params, self._custom_button_callbacks)
 
     if not self._setup:
       if self._alerts_layout.active_alerts() > 0:

openpilot/selfdrive/ui/sunnypilot/custom_button.py

@@ -9,15 +9,11 @@ class CustomButtonAction(IntEnum):
   CYCLE_UI = 3
 
 
-def handle_custom_button(sm, params, callbacks):
-  if not sm.updated['carState']:
-    return
-
-  custom_pressed = any(be.type == car.CarState.ButtonEvent.Type.altButton2 and be.pressed
-                       for be in sm['carState'].buttonEvents)
-  if not custom_pressed:
-    return
-
-  action = CustomButtonAction(params.get('CustomButtonAction', return_default=True))
-  if callback := callbacks.get(action):
-    callback()
+def handle_custom_button(messages, params, callbacks):
+  for msg in messages:
+    custom_pressed = any(be.type == car.CarState.ButtonEvent.Type.altButton2 and be.pressed
+                         for be in msg.carState.buttonEvents)
+    if custom_pressed:
+      action = CustomButtonAction(params.get('CustomButtonAction', return_default=True))
+      if callback := callbacks.get(action):
+        callback()

openpilot/selfdrive/ui/tests/test_custom_button.py

@@ -6,26 +6,16 @@ from opendbc.car.structs import car
 from openpilot.selfdrive.ui.sunnypilot.custom_button import CustomButtonAction, handle_custom_button
 
 
-class FakeSubMaster:
-  def __init__(self, messages):
-    self.messages = messages
-    self.updated = {'carState': True}
-
-  def __getitem__(self, key):
-    return self.messages[key]
-
-
 def test_custom_button_actions():
   params = Mock()
-  sm = FakeSubMaster({
-    'carState': SimpleNamespace(buttonEvents=[SimpleNamespace(
+  press = SimpleNamespace(carState=SimpleNamespace(buttonEvents=[SimpleNamespace(
       type=car.CarState.ButtonEvent.Type.altButton2,
       pressed=True,
-    )]),
-  })
+    )]))
+  messages = [press, SimpleNamespace(carState=SimpleNamespace(buttonEvents=[])), press]
   callbacks = {action: Mock() for action in CustomButtonAction if action != CustomButtonAction.NONE}
 
   for action, callback in callbacks.items():
     params.get.return_value = action
-    handle_custom_button(sm, params, callbacks)
-    callback.assert_called_once()
+    handle_custom_button(messages, params, callbacks)
+    assert callback.call_count == 2