feat(core): unify RustLayout, implement single global layout

common/protob/messages-debug.proto

@@ -51,7 +51,7 @@ message DebugLinkDecision {
 
     optional uint32 x = 4;                             // touch X coordinate
     optional uint32 y = 5;                             // touch Y coordinate
-    optional bool wait = 6;                            // wait for layout change
+    optional bool wait = 6 [deprecated=true];          // wait for layout change
     optional uint32 hold_ms = 7;                       // touch hold duration
     optional DebugPhysicalButton physical_button = 8;  // physical button press
 }
@@ -61,6 +61,7 @@ message DebugLinkDecision {
  * @end
  */
 message DebugLinkLayout {
+    option deprecated = true;
     repeated string tokens = 1;
 }
 
@@ -89,9 +90,26 @@ message DebugLinkRecordScreen {
  * @next DebugLinkState
  */
 message DebugLinkGetState {
-    optional bool wait_word_list = 1;  // Trezor T only - wait until mnemonic words are shown
+    /// Wait behavior of the call.
-    optional bool wait_word_pos = 2;   // Trezor T only - wait until reset word position is requested
+    enum DebugWaitType {
-    optional bool wait_layout = 3;     // wait until current layout changes
+        /// Respond immediately. If no layout is currently displayed, the layout
+        /// response will be empty.
+        IMMEDIATE = 0;
+        /// Wait for next layout. If a layout is displayed, waits for it to change.
+        /// If no layout is displayed, waits for one to come up.
+        NEXT_LAYOUT = 1;
+        /// Return current layout. If no layout is currently displayed, waits for
+        /// one to come up.
+        CURRENT_LAYOUT = 2;
+    }
+
+    // Trezor T < 2.6.0 only - wait until mnemonic words are shown
+    optional bool wait_word_list = 1 [deprecated=true];
+    // Trezor T < 2.6.0 only - wait until reset word position is requested
+    optional bool wait_word_pos = 2  [deprecated=true];
+    // trezor-core only - wait until current layout changes
+    // changed in 2.6.4: multiple wait types instead of true/false.
+    optional DebugWaitType wait_layout = 3 [default=IMMEDIATE];
 }
 
 /**
@@ -192,6 +210,7 @@ message DebugLinkEraseSdCard {
  * @next Success
  */
 message DebugLinkWatchLayout {
+    option deprecated = true;
     optional bool watch = 1;  // if true, start watching layout.
                               // if false, stop.
 }
@@ -203,4 +222,5 @@ message DebugLinkWatchLayout {
  * @next Success
  */
 message DebugLinkResetDebugEvents {
+    option deprecated = true;
 }

core/src/all_modules.py

@@ -105,6 +105,8 @@ trezor.enums.DebugPhysicalButton
 import trezor.enums.DebugPhysicalButton
 trezor.enums.DebugSwipeDirection
 import trezor.enums.DebugSwipeDirection
+trezor.enums.DebugWaitType
+import trezor.enums.DebugWaitType
 trezor.enums.DecredStakingSpendType
 import trezor.enums.DecredStakingSpendType
 trezor.enums.FailureType
@@ -167,10 +169,6 @@ trezor.ui.layouts.tr
 import trezor.ui.layouts.tr
 trezor.ui.layouts.tr.fido
 import trezor.ui.layouts.tr.fido
-trezor.ui.layouts.tr.homescreen
-import trezor.ui.layouts.tr.homescreen
-trezor.ui.layouts.tr.progress
-import trezor.ui.layouts.tr.progress
 trezor.ui.layouts.tr.recovery
 import trezor.ui.layouts.tr.recovery
 trezor.ui.layouts.tr.reset
@@ -179,10 +177,6 @@ trezor.ui.layouts.tt
 import trezor.ui.layouts.tt
 trezor.ui.layouts.tt.fido
 import trezor.ui.layouts.tt.fido
-trezor.ui.layouts.tt.homescreen
-import trezor.ui.layouts.tt.homescreen
-trezor.ui.layouts.tt.progress
-import trezor.ui.layouts.tt.progress
 trezor.ui.layouts.tt.recovery
 import trezor.ui.layouts.tt.recovery
 trezor.ui.layouts.tt.reset

core/src/apps/bitcoin/sign_tx/layout.py

@@ -21,7 +21,6 @@ if TYPE_CHECKING:
 
     from trezor.enums import AmountUnit
     from trezor.messages import TxAckPaymentRequest, TxOutput
-    from trezor.ui.layouts import LayoutType
 
     from apps.common.coininfo import CoinInfo
     from apps.common.paths import Bip32Path
@@ -73,7 +72,7 @@ async def confirm_output(
         assert data is not None
         if omni.is_valid(data):
             # OMNI transaction
-            layout: LayoutType = confirm_metadata(
+            layout = confirm_metadata(
                 "omni_transaction",
                 "OMNI transaction",
                 omni.parse(data),

core/src/apps/bitcoin/sign_tx/progress.py

@@ -117,14 +117,10 @@ class Progress:
 
         progress_layout = coinjoin_progress if self.is_coinjoin else bitcoin_progress
         workflow.close_others()
-        text = "Signing transaction..." if self.signing else "Loading transaction..."
+        text = "Signing transaction" if self.signing else "Loading transaction"
         self.progress_layout = progress_layout(text)
 
     def report(self) -> None:
-        from trezor import utils
-
-        if utils.DISABLE_ANIMATION:
-            return
         p = int(1000 * self.progress / self.steps)
         self.progress_layout.report(p)
 

core/src/apps/common/mnemonic.py

@@ -5,7 +5,7 @@ from trezor import utils
 
 if TYPE_CHECKING:
     from trezor.enums import BackupType
-    from trezor.ui.layouts.common import ProgressLayout
+    from trezor.ui import ProgressLayout
 
 
 def get() -> tuple[bytes | None, BackupType]:

core/src/apps/debug/__init__.py

@@ -8,49 +8,36 @@ if __debug__:
 
     import trezorui2
     from storage import debug as storage
-    from storage.debug import debug_events
+    from trezor import io, log, loop, ui, utils, wire, workflow
-    from trezor import log, loop, utils, wire
+    from trezor.enums import DebugWaitType, MessageType
-    from trezor.enums import MessageType
+    from trezor.messages import Success
-    from trezor.messages import DebugLinkLayout, Success
     from trezor.ui import display
-    from trezor.wire import context
-
-    from apps import workflow_handlers
 
     if TYPE_CHECKING:
+        from typing import Any, Awaitable, Callable
+
+        from trezor.enums import DebugButton, DebugPhysicalButton, DebugSwipeDirection
         from trezor.messages import (
             DebugLinkDecision,
             DebugLinkEraseSdCard,
             DebugLinkGetState,
             DebugLinkRecordScreen,
             DebugLinkReseedRandom,
-            DebugLinkResetDebugEvents,
             DebugLinkState,
-            DebugLinkWatchLayout,
         )
         from trezor.ui import Layout
+        from trezor.wire import WireInterface, context
 
-    swipe_chan = loop.chan()
+        Handler = Callable[[Any], Awaitable[Any]]
-    result_chan = loop.chan()
-    button_chan = loop.chan()
-    click_chan = loop.chan()
-    swipe_signal = swipe_chan.take
-    result_signal = result_chan.take
-    button_signal = button_chan.take
-    click_signal = click_chan.take
 
-    debuglink_decision_chan = loop.chan()
+    layout_change_chan = loop.mailbox()
-
-    layout_change_chan = loop.chan()
 
     DEBUG_CONTEXT: context.Context | None = None
 
-    LAYOUT_WATCHER_NONE = 0
-    LAYOUT_WATCHER_STATE = 1
-    LAYOUT_WATCHER_LAYOUT = 2
-
     REFRESH_INDEX = 0
 
+    _DEADLOCK_DETECT_SLEEP = loop.sleep(2000)
+
     def screenshot() -> bool:
         if storage.save_screen:
             # Starting with "refresh00", allowing for 100 emulator restarts
@@ -61,158 +48,239 @@ if __debug__:
             return True
         return False
 
-    def notify_layout_change(layout: Layout, event_id: int | None = None) -> None:
+    def notify_layout_change(layout: Layout | None) -> None:
-        layout.read_content_into(storage.current_content_tokens)
+        layout_change_chan.put(layout, replace=True)
-        if storage.watch_layout_changes or layout_change_chan.takers:
-            payload = (event_id, storage.current_content_tokens)
-            layout_change_chan.publish(payload)
 
-    async def _dispatch_debuglink_decision(
+    def wait_until_layout_is_running(limit: int | None = None) -> Awaitable[None]:  # type: ignore [awaitable-is-generator]
-        event_id: int | None, msg: DebugLinkDecision
+        counter = 0
+        while ui.CURRENT_LAYOUT is None:
+            yield
+            if limit is not None and counter > limit:
+                return
+
+    async def return_layout_change(
+        ctx: wire.context.Context, detect_deadlock: bool = False
     ) -> None:
-        from trezor.enums import DebugButton
+        # set up the wait
+        storage.layout_watcher = True
 
-        if msg.button is not None:
+        # wait for layout change
-            if msg.button == DebugButton.NO:
+        while True:
-                await result_chan.put((event_id, trezorui2.CANCELLED))
+            if not detect_deadlock or not layout_change_chan.is_empty():
-            elif msg.button == DebugButton.YES:
+                # short-circuit if there is a result already waiting
-                await result_chan.put((event_id, trezorui2.CONFIRMED))
+                next_layout = await layout_change_chan
-            elif msg.button == DebugButton.INFO:
-                await result_chan.put((event_id, trezorui2.INFO))
             else:
-                raise RuntimeError(f"Invalid msg.button - {msg.button}")
+                next_layout = await loop.race(
-        elif msg.input is not None:
+                    layout_change_chan, _DEADLOCK_DETECT_SLEEP
-            await result_chan.put((event_id, msg.input))
-        elif msg.swipe is not None:
-            await swipe_chan.put((event_id, msg.swipe))
-        else:
-            # Sanity check. The message will be visible in terminal.
-            raise RuntimeError("Invalid DebugLinkDecision message")
-
-    async def debuglink_decision_dispatcher() -> None:
-        while True:
-            event_id, msg = await debuglink_decision_chan.take()
-            await _dispatch_debuglink_decision(event_id, msg)
-
-    async def get_layout_change_content() -> list[str]:
-        awaited_event_id = debug_events.awaited_event
-        last_result_id = debug_events.last_result
-
-        if awaited_event_id is not None and awaited_event_id == last_result_id:
-            # We are awaiting the event that just happened - return current state
-            return storage.current_content_tokens
-
-        while True:
-            event_id, content = await layout_change_chan.take()
-            if awaited_event_id is None or event_id is None:
-                # Not waiting for anything or event does not have ID
-                break
-            elif event_id == awaited_event_id:
-                # We found what we were waiting for
-                debug_events.awaited_event = None
-                break
-            elif event_id > awaited_event_id:
-                # Sanity check
-                raise RuntimeError(
-                    f"Waiting for event that already happened - {event_id} > {awaited_event_id}"
                 )
 
-        if awaited_event_id is not None:
+            if next_layout is None:
-            # Updating last result
+                # layout close event. loop again
-            debug_events.last_result = awaited_event_id
+                continue
 
-        return content
+            if isinstance(next_layout, ui.Layout):
+                break
 
-    async def return_layout_change() -> None:
+            if isinstance(next_layout, int):
-        content_tokens = await get_layout_change_content()
+                # sleep result from the deadlock detector
+                raise wire.FirmwareError("layout deadlock detected")
 
-        assert DEBUG_CONTEXT is not None
+            raise RuntimeError(
-        if storage.layout_watcher is LAYOUT_WATCHER_LAYOUT:
+                f"Unexpected layout change: {next_layout}, {type(next_layout)}"
-            await DEBUG_CONTEXT.write(DebugLinkLayout(tokens=content_tokens))
+            )
+
+        assert ui.CURRENT_LAYOUT is next_layout
+
+        # send the message and reset the wait
+        storage.layout_watcher = False
+        await ctx.write(_state())
+
+    async def _layout_click(layout: Layout, x: int, y: int, hold_ms: int = 0) -> None:
+        msg = layout.layout.touch_event(io.TOUCH_START, x, y)
+        layout._emit_message(msg)
+        layout._paint()
+
+        if hold_ms:
+            await loop.sleep(hold_ms)
+            workflow.idle_timer.touch()
+
+        msg = layout.layout.touch_event(io.TOUCH_END, x, y)
+        layout._emit_message(msg)
+        layout._paint()
+
+    async def _layout_press_button(
+        layout: Layout, debug_btn: DebugPhysicalButton, hold_ms: int = 0
+    ) -> None:
+        from trezor.enums import DebugPhysicalButton
+
+        buttons = []
+
+        if debug_btn == DebugPhysicalButton.LEFT_BTN:
+            buttons.append(io.BUTTON_LEFT)
+        elif debug_btn == DebugPhysicalButton.RIGHT_BTN:
+            buttons.append(io.BUTTON_RIGHT)
+        elif debug_btn == DebugPhysicalButton.MIDDLE_BTN:
+            buttons.append(io.BUTTON_LEFT)
+            buttons.append(io.BUTTON_RIGHT)
+
+        for btn in buttons:
+            msg = layout.layout.button_event(io.BUTTON_PRESSED, btn)
+            layout._emit_message(msg)
+            layout._paint()
+
+        if hold_ms:
+            await loop.sleep(hold_ms)
+            workflow.idle_timer.touch()
+
+        for btn in buttons:
+            msg = layout.layout.button_event(io.BUTTON_RELEASED, btn)
+            layout._emit_message(msg)
+            layout._paint()
+
+    if utils.USE_TOUCH:
+
+        async def _layout_swipe(layout: Layout, direction: DebugSwipeDirection) -> None:  # type: ignore [obscured by a declaration of the same name]
+            from trezor.enums import DebugSwipeDirection
+
+            orig_x = orig_y = 120
+            off_x, off_y = {
+                DebugSwipeDirection.UP: (0, -30),
+                DebugSwipeDirection.DOWN: (0, 30),
+                DebugSwipeDirection.LEFT: (-30, 0),
+                DebugSwipeDirection.RIGHT: (30, 0),
+            }[direction]
+
+            for event, x, y in (
+                (io.TOUCH_START, orig_x, orig_y),
+                (io.TOUCH_MOVE, orig_x + 1 * off_x, orig_y + 1 * off_y),
+                (io.TOUCH_END, orig_x + 2 * off_x, orig_y + 2 * off_y),
+            ):
+                msg = layout.layout.touch_event(event, x, y)
+                layout._emit_message(msg)
+                layout._paint()
+
+    elif utils.USE_BUTTON:
+
+        def _layout_swipe(
+            layout: Layout, direction: DebugSwipeDirection
+        ) -> Awaitable[None]:
+            from trezor.enums import DebugPhysicalButton, DebugSwipeDirection
+
+            if direction == DebugSwipeDirection.UP:
+                button = DebugPhysicalButton.RIGHT_BTN
+            elif direction == DebugSwipeDirection.DOWN:
+                button = DebugPhysicalButton.LEFT_BTN
+            else:
+                raise RuntimeError  # unsupported swipe direction on TR
+
+            return _layout_press_button(layout, button)
+
+    else:
+        raise RuntimeError  # No way to swipe with no buttons and no touches
+
+    async def _layout_event(layout: Layout, button: DebugButton) -> None:
+        from trezor.enums import DebugButton
+
+        if button == DebugButton.NO:
+            layout._emit_message(trezorui2.CANCELLED)
+        elif button == DebugButton.YES:
+            layout._emit_message(trezorui2.CONFIRMED)
+        elif button == DebugButton.INFO:
+            layout._emit_message(trezorui2.INFO)
         else:
-            from trezor.messages import DebugLinkState
+            raise RuntimeError("Invalid DebugButton")
 
-            await DEBUG_CONTEXT.write(DebugLinkState(tokens=content_tokens))
+    async def dispatch_DebugLinkDecision(
-        storage.layout_watcher = LAYOUT_WATCHER_NONE
+        msg: DebugLinkDecision,
-
+    ) -> DebugLinkState | None:
-    async def dispatch_DebugLinkWatchLayout(msg: DebugLinkWatchLayout) -> Success:
+        from trezor import ui, workflow
-        from trezor import ui
-
-        layout_change_chan.putters.clear()
-        if msg.watch:
-            await ui.wait_until_layout_is_running()
-        storage.watch_layout_changes = bool(msg.watch)
-        log.debug(__name__, "Watch layout changes: %s", storage.watch_layout_changes)
-        return Success()
-
-    async def dispatch_DebugLinkResetDebugEvents(
-        msg: DebugLinkResetDebugEvents,
-    ) -> Success:
-        # Resetting the debug events makes sure that the previous
-        # events/layouts are not mixed with the new ones.
-        storage.reset_debug_events()
-        return Success()
-
-    async def dispatch_DebugLinkDecision(msg: DebugLinkDecision) -> None:
-        from trezor import workflow
 
         workflow.idle_timer.touch()
 
-        if debuglink_decision_chan.putters:
-            log.warning(__name__, "DebugLinkDecision queue is not empty")
-
         x = msg.x  # local_cache_attribute
         y = msg.y  # local_cache_attribute
 
-        # Incrementing the counter for last events so we know what to await
+        await wait_until_layout_is_running()
-        debug_events.last_event += 1
+        layout = ui.CURRENT_LAYOUT
+        assert layout is not None
+        assert isinstance(layout, ui.Layout)
+        layout_change_chan.clear()
 
-        # click on specific coordinates, with possible hold
+        try:
-        if x is not None and y is not None:
+            # click on specific coordinates, with possible hold
-            click_chan.publish((debug_events.last_event, x, y, msg.hold_ms))
+            if x is not None and y is not None:
-        # press specific button
+                await _layout_click(layout, x, y, msg.hold_ms or 0)
-        elif msg.physical_button is not None:
+            # press specific button
-            button_chan.publish(
+            elif msg.physical_button is not None:
-                (debug_events.last_event, msg.physical_button, msg.hold_ms)
+                await _layout_press_button(
-            )
+                    layout, msg.physical_button, msg.hold_ms or 0
-        else:
+                )
-            # Will get picked up by _dispatch_debuglink_decision eventually
+            elif msg.swipe is not None:
-            debuglink_decision_chan.publish((debug_events.last_event, msg))
+                await _layout_swipe(layout, msg.swipe)
+            elif msg.button is not None:
+                await _layout_event(layout, msg.button)
+            elif msg.input is not None:
+                layout._emit_message(msg.input)
+            else:
+                raise RuntimeError("Invalid DebugLinkDecision message")
 
-        if msg.wait:
+        except ui.Shutdown:
-            # We wait for all the previously sent events
+            # Shutdown should be raised if the layout is supposed to stop after
-            debug_events.awaited_event = debug_events.last_event
+            # processing the event. In that case, we need to yield to give the layout
-            storage.layout_watcher = LAYOUT_WATCHER_LAYOUT
+            # callers time to finish their jobs. We want to make sure that the handling
-            loop.schedule(return_layout_change())
+            # does not continue until the event is truly processed.
+            result = await layout_change_chan
+            assert result is None
 
-    async def dispatch_DebugLinkGetState(
+        # If no exception was raised, the layout did not shut down. That means that it
-        msg: DebugLinkGetState,
+        # just updated itself. The update is already live for the caller to retrieve.
-    ) -> DebugLinkState | None:
+
+    def _state() -> DebugLinkState:
         from trezor.messages import DebugLinkState
 
         from apps.common import mnemonic, passphrase
 
-        m = DebugLinkState()
+        tokens = []
-        m.mnemonic_secret = mnemonic.get_secret()
-        m.mnemonic_type = mnemonic.get_type()
-        m.passphrase_protection = passphrase.is_enabled()
-        m.reset_entropy = storage.reset_internal_entropy
 
-        if msg.wait_layout:
+        def callback(*args: str) -> None:
-            if not storage.watch_layout_changes:
+            tokens.extend(args)
-                raise wire.ProcessError("Layout is not watched")
+
-            storage.layout_watcher = LAYOUT_WATCHER_STATE
+        if ui.CURRENT_LAYOUT is not None:
-            # We wait for the last previously sent event to finish
+            ui.CURRENT_LAYOUT.layout.trace(callback)
-            debug_events.awaited_event = debug_events.last_event
+
-            loop.schedule(return_layout_change())
+        print("!!! reporting state:", "".join(tokens))
-            return None
+
+        return DebugLinkState(
+            mnemonic_secret=mnemonic.get_secret(),
+            mnemonic_type=mnemonic.get_type(),
+            passphrase_protection=passphrase.is_enabled(),
+            reset_entropy=storage.reset_internal_entropy,
+            tokens=tokens,
+        )
+
+    async def dispatch_DebugLinkGetState(
+        msg: DebugLinkGetState,
+    ) -> DebugLinkState | None:
+        if msg.wait_layout == DebugWaitType.IMMEDIATE:
+            return _state()
+
+        assert DEBUG_CONTEXT is not None
+        if msg.wait_layout == DebugWaitType.NEXT_LAYOUT:
+            layout_change_chan.clear()
+            return await return_layout_change(DEBUG_CONTEXT, detect_deadlock=False)
+
+        # default behavior: msg.wait_layout == DebugWaitType.CURRENT_LAYOUT
+        if not isinstance(ui.CURRENT_LAYOUT, ui.Layout):
+            return await return_layout_change(DEBUG_CONTEXT, detect_deadlock=True)
         else:
-            m.tokens = storage.current_content_tokens
+            return _state()
-
-        return m
 
     async def dispatch_DebugLinkRecordScreen(msg: DebugLinkRecordScreen) -> Success:
         if msg.target_directory:
+            # Ensure we consistently start at a layout, instead of randomly sometimes
+            # hitting the pause between layouts and rendering the "upcoming" one.
+            await wait_until_layout_is_running()
+
             # In case emulator is restarted but we still want to record screenshots
             # into the same directory as before, we need to increment the refresh index,
             # so that the screenshots are not overwritten.
@@ -220,6 +288,10 @@ if __debug__:
             REFRESH_INDEX = msg.refresh_index
             storage.save_screen_directory = msg.target_directory
             storage.save_screen = True
+
+            # invoke the refresh function to save the first freshly painted screenshot.
+            display.refresh()
+
         else:
             storage.save_screen = False
             display.clear_save()  # clear C buffers
@@ -255,19 +327,87 @@ if __debug__:
             sdcard.power_off()
         return Success()
 
+    async def _no_op(msg: Any) -> Success:
+        return Success()
+
+    WIRE_BUFFER_DEBUG = bytearray(1024)
+
+    async def handle_session(iface: WireInterface) -> None:
+        from trezor import protobuf, wire
+        from trezor.wire import codec_v1, context
+
+        global DEBUG_CONTEXT
+
+        DEBUG_CONTEXT = ctx = context.Context(iface, 0, WIRE_BUFFER_DEBUG)
+
+        if storage.layout_watcher:
+            try:
+                await return_layout_change(ctx)
+            except Exception as e:
+                log.exception(__name__, e)
+
+        while True:
+            try:
+                try:
+                    msg = await ctx.read_from_wire()
+                except codec_v1.CodecError as exc:
+                    log.exception(__name__, exc)
+                    await ctx.write(wire.failure(exc))
+                    continue
+
+                req_type = None
+                try:
+                    req_type = protobuf.type_for_wire(msg.type)
+                    msg_type = req_type.MESSAGE_NAME
+                except Exception:
+                    msg_type = f"{msg.type} - unknown message type"
+                log.debug(
+                    __name__,
+                    "%s:%x receive: <%s>",
+                    ctx.iface.iface_num(),
+                    ctx.sid,
+                    msg_type,
+                )
+
+                if msg.type not in WORKFLOW_HANDLERS:
+                    await ctx.write(wire.unexpected_message())
+                    continue
+
+                elif req_type is None:
+                    # Message type is in workflow handlers but not in protobuf
+                    # definitions. This indicates a deprecated message.
+                    # We put a no-op handler for those messages.
+                    # XXX return a Failure here?
+                    await ctx.write(Success())
+                    continue
+
+                req_msg = wire.wrap_protobuf_load(msg.data, req_type)
+                try:
+                    res_msg = await WORKFLOW_HANDLERS[msg.type](req_msg)
+                except Exception as exc:
+                    # Log and ignore, never die.
+                    log.exception(__name__, exc)
+                    res_msg = wire.failure(exc)
+
+                if res_msg is not None:
+                    await ctx.write(res_msg)
+
+            except Exception as exc:
+                # Log and try again. This should only happen for USB errors and we
+                # try to stay robust in such case.
+                log.exception(__name__, exc)
+
+    WORKFLOW_HANDLERS: dict[int, Handler] = {
+        MessageType.DebugLinkDecision: dispatch_DebugLinkDecision,
+        MessageType.DebugLinkGetState: dispatch_DebugLinkGetState,
+        MessageType.DebugLinkReseedRandom: dispatch_DebugLinkReseedRandom,
+        MessageType.DebugLinkRecordScreen: dispatch_DebugLinkRecordScreen,
+        MessageType.DebugLinkEraseSdCard: dispatch_DebugLinkEraseSdCard,
+        MessageType.DebugLinkWatchLayout: _no_op,
+        MessageType.DebugLinkResetDebugEvents: _no_op,
+    }
+
     def boot() -> None:
-        register = workflow_handlers.register  # local_cache_attribute
+        import usb
 
-        register(MessageType.DebugLinkDecision, dispatch_DebugLinkDecision)  # type: ignore [Argument of type "(msg: DebugLinkDecision) -> Coroutine[Any, Any, None]" cannot be assigned to parameter "handler" of type "Handler[Msg@register]" in function "register"]
+        loop.schedule(handle_session(usb.iface_debug))
-        register(MessageType.DebugLinkGetState, dispatch_DebugLinkGetState)  # type: ignore [Argument of type "(msg: DebugLinkGetState) -> Coroutine[Any, Any, DebugLinkState | None]" cannot be assigned to parameter "handler" of type "Handler[Msg@register]" in function "register"]
-        register(MessageType.DebugLinkReseedRandom, dispatch_DebugLinkReseedRandom)
-        register(MessageType.DebugLinkRecordScreen, dispatch_DebugLinkRecordScreen)
-        register(MessageType.DebugLinkEraseSdCard, dispatch_DebugLinkEraseSdCard)
-        register(MessageType.DebugLinkWatchLayout, dispatch_DebugLinkWatchLayout)
-        register(
-            MessageType.DebugLinkResetDebugEvents, dispatch_DebugLinkResetDebugEvents
-        )
-
-        loop.schedule(debuglink_decision_dispatcher())
-        if storage.layout_watcher is not LAYOUT_WATCHER_NONE:
-            loop.schedule(return_layout_change())

core/src/apps/homescreen/__init__.py

@@ -12,7 +12,7 @@ from apps.common.authorization import is_set_any_session
 
 
 async def busyscreen() -> None:
-    await Busyscreen(busy_expiry_ms())
+    await Busyscreen(busy_expiry_ms()).get_result()
 
 
 async def homescreen() -> None:
@@ -43,7 +43,7 @@ async def homescreen() -> None:
         notification=notification,
         notification_is_error=notification_is_error,
         hold_to_lock=config.has_pin(),
-    )
+    ).get_result()
     lock_device()
 
 
@@ -57,7 +57,7 @@ async def _lockscreen(screensaver: bool = False) -> None:
         await Lockscreen(
             label=storage.device.get_label(),
             coinjoin_authorized=is_set_any_session(MessageType.AuthorizeCoinJoin),
-        )
+        ).get_result()
     # Otherwise proceed directly to unlock() call. If the device is already unlocked,
     # it should be a no-op storage-wise, but it resets the internal configuration
     # to an unlocked state.

core/src/apps/management/recovery_device/layout.py

@@ -41,18 +41,20 @@ async def _confirm_abort(dry_run: bool = False) -> None:
 async def request_mnemonic(
     word_count: int, backup_type: BackupType | None
 ) -> str | None:
-    from trezor.ui.layouts.common import button_request
     from trezor.ui.layouts.recovery import request_word
 
     from . import word_validity
 
-    await button_request("mnemonic", code=ButtonRequestType.MnemonicInput)
-
     words: list[str] = []
+    send_button_request = True
     for i in range(word_count):
         word = await request_word(
-            i, word_count, is_slip39=backup_types.is_slip39_word_count(word_count)
+            i,
+            word_count,
+            backup_types.is_slip39_word_count(word_count),
+            send_button_request,
         )
+        send_button_request = False
         words.append(word)
 
         try:

core/src/apps/management/reset_device/layout.py

@@ -128,7 +128,7 @@ async def _show_confirmation_success(
             subheader = f"Group {group_index + 1} - Share {share_index + 1} checked successfully."
             text = "Continue with the next share."
 
-    return await show_success("success_recovery", text, subheader)
+    await show_success("success_recovery", text, subheader)
 
 
 async def _show_confirmation_failure() -> None:

core/src/apps/misc/get_firmware_hash.py

@@ -2,7 +2,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from trezor.messages import FirmwareHash, GetFirmwareHash
-    from trezor.ui.layouts.common import ProgressLayout
+    from trezor.ui import ProgressLayout
 
 _progress_obj: ProgressLayout | None = None
 
@@ -15,7 +15,7 @@ async def get_firmware_hash(msg: GetFirmwareHash) -> FirmwareHash:
 
     workflow.close_others()
     global _progress_obj
-    _progress_obj = progress()
+    _progress_obj = progress("PLEASE WAIT", "")
 
     try:
         hash = firmware_hash(msg.challenge, _render_progress)

core/src/apps/webauthn/fido2.py

@@ -8,7 +8,8 @@ import storage.device as storage_device
 from trezor import config, io, log, loop, utils, wire, workflow
 from trezor.crypto import hashlib
 from trezor.crypto.curve import nist256p1
-from trezor.ui.layouts import show_error_popup
+from trezor.ui import Layout
+from trezor.ui.layouts import error_popup
 
 from apps.base import set_homescreen
 from apps.common import cbor
@@ -615,16 +616,36 @@ async def _confirm_fido(title: str, credential: Credential) -> bool:
         return False
 
 
+async def _show_error_popup(
+    title: str,
+    description: str,
+    subtitle: str | None = None,
+    description_param: str = "",
+    *,
+    button: str = "",
+    timeout_ms: int = 0,
+) -> None:
+    popup = error_popup(
+        title,
+        description,
+        subtitle,
+        description_param,
+        button=button,
+        timeout_ms=timeout_ms,
+    )
+    await Layout(popup).get_result()
+
+
 async def _confirm_bogus_app(title: str) -> None:
     if _last_auth_valid:
-        await show_error_popup(
+        await _show_error_popup(
             title,
             "This device is already registered with this application.",
             "Already registered.",
             timeout_ms=_POPUP_TIMEOUT_MS,
         )
     else:
-        await show_error_popup(
+        await _show_error_popup(
             title,
             "This device is not registered with this application.",
             "Not registered.",
@@ -841,7 +862,7 @@ class Fido2ConfirmExcluded(Fido2ConfirmMakeCredential):
         await send_cmd(cmd, self.iface)
         self.finished = True
 
-        await show_error_popup(
+        await _show_error_popup(
             "FIDO2 Register",
             "This device is already registered with {}.",
             "Already registered.",
@@ -924,7 +945,7 @@ class Fido2ConfirmNoPin(State):
         await send_cmd(cmd, self.iface)
         self.finished = True
 
-        await show_error_popup(
+        await _show_error_popup(
             "FIDO2 Verify User",
             "Please enable PIN protection.",
             "Unable to verify user.",
@@ -947,7 +968,7 @@ class Fido2ConfirmNoCredentials(Fido2ConfirmGetAssertion):
         await send_cmd(cmd, self.iface)
         self.finished = True
 
-        await show_error_popup(
+        await _show_error_popup(
             "FIDO2 Authenticate",
             "This device is not registered with\n{}.",
             "Not registered.",
@@ -1056,6 +1077,7 @@ class DialogManager:
 
         try:
             while self.result is _RESULT_NONE:
+                workflow.close_others()
                 result = await self.state.confirm_dialog()
                 if isinstance(result, State):
                     self.state = result

core/src/boot.py

@@ -34,7 +34,7 @@ async def bootscreen() -> None:
     while True:
         try:
             if can_lock_device():
-                await lockscreen
+                await lockscreen.get_result()
             await verify_user_pin()
             storage.init_unlocked()
             allow_all_loader_messages()

core/src/storage/debug.py

@@ -7,28 +7,6 @@ if __debug__:
     save_screen = False
     save_screen_directory = "."
 
-    current_content_tokens: list[str] = [""] * 60
+    layout_watcher = False
-    current_content_tokens.clear()
-
-    watch_layout_changes = False
-    layout_watcher = 0
 
     reset_internal_entropy: bytes = b""
-
-    class DebugEvents:
-        def __init__(self):
-            self.reset()
-
-        def reset(self) -> None:
-            self.last_event = 0
-            self.last_result: int | None = None
-            self.awaited_event: int | None = None
-
-    debug_events = DebugEvents()
-
-    def reset_debug_events() -> None:
-        debug_events.reset()
-
-    # Event resulted in the layout change, call
-    # notify_layout_change with this ID in first_paint of next layout.
-    new_layout_event_id: int | None = None

core/src/trezor/enums/DebugWaitType.py

@@ -0,0 +1,7 @@
+# Automatically generated by pb2py
+# fmt: off
+# isort:skip_file
+
+IMMEDIATE = 0
+NEXT_LAYOUT = 1
+CURRENT_LAYOUT = 2

core/src/trezor/enums/__init__.py

@@ -470,6 +470,11 @@ if TYPE_CHECKING:
         MIDDLE_BTN = 1
         RIGHT_BTN = 2
 
+    class DebugWaitType(IntEnum):
+        IMMEDIATE = 0
+        NEXT_LAYOUT = 1
+        CURRENT_LAYOUT = 2
+
     class EthereumDefinitionType(IntEnum):
         NETWORK = 0
         TOKEN = 1

core/src/trezor/loop.py

@@ -666,24 +666,3 @@ class spawn(Syscall):
         is True, it would be calling close on self, which will result in a ValueError.
         """
         return self.task is this_task
-
-
-class Timer(Syscall):
-    def __init__(self) -> None:
-        self.task: Task | None = None
-        # Event::Attach is evaluated before task is set. Use this list to
-        # buffer timers until task is set.
-        self.before_task: list[tuple[int, Any]] = []
-
-    def handle(self, task: Task) -> None:
-        self.task = task
-        for deadline, value in self.before_task:
-            schedule(self.task, value, deadline)
-        self.before_task.clear()
-
-    def schedule(self, deadline: int, value: Any) -> None:
-        deadline = utime.ticks_add(utime.ticks_ms(), deadline)
-        if self.task is not None:
-            schedule(self.task, value, deadline)
-        else:
-            self.before_task.append((deadline, value))

core/src/trezor/messages.py

@@ -38,6 +38,7 @@ if TYPE_CHECKING:
     from trezor.enums import DebugButton  # noqa: F401
     from trezor.enums import DebugPhysicalButton  # noqa: F401
     from trezor.enums import DebugSwipeDirection  # noqa: F401
+    from trezor.enums import DebugWaitType  # noqa: F401
     from trezor.enums import DecredStakingSpendType  # noqa: F401
     from trezor.enums import EthereumDataType  # noqa: F401
     from trezor.enums import EthereumDefinitionType  # noqa: F401
@@ -2701,7 +2702,6 @@ if TYPE_CHECKING:
         input: "str | None"
         x: "int | None"
         y: "int | None"
-        wait: "bool | None"
         hold_ms: "int | None"
         physical_button: "DebugPhysicalButton | None"
 
@@ -2713,7 +2713,6 @@ if TYPE_CHECKING:
             input: "str | None" = None,
             x: "int | None" = None,
             y: "int | None" = None,
-            wait: "bool | None" = None,
             hold_ms: "int | None" = None,
             physical_button: "DebugPhysicalButton | None" = None,
         ) -> None:
@@ -2723,20 +2722,6 @@ if TYPE_CHECKING:
         def is_type_of(cls, msg: Any) -> TypeGuard["DebugLinkDecision"]:
             return isinstance(msg, cls)
 
-    class DebugLinkLayout(protobuf.MessageType):
-        tokens: "list[str]"
-
-        def __init__(
-            self,
-            *,
-            tokens: "list[str] | None" = None,
-        ) -> None:
-            pass
-
-        @classmethod
-        def is_type_of(cls, msg: Any) -> TypeGuard["DebugLinkLayout"]:
-            return isinstance(msg, cls)
-
     class DebugLinkReseedRandom(protobuf.MessageType):
         value: "int | None"
 
@@ -2768,16 +2753,12 @@ if TYPE_CHECKING:
             return isinstance(msg, cls)
 
     class DebugLinkGetState(protobuf.MessageType):
-        wait_word_list: "bool | None"
+        wait_layout: "DebugWaitType"
-        wait_word_pos: "bool | None"
-        wait_layout: "bool | None"
 
         def __init__(
             self,
             *,
-            wait_word_list: "bool | None" = None,
+            wait_layout: "DebugWaitType | None" = None,
-            wait_word_pos: "bool | None" = None,
-            wait_layout: "bool | None" = None,
         ) -> None:
             pass
 
@@ -2923,26 +2904,6 @@ if TYPE_CHECKING:
         def is_type_of(cls, msg: Any) -> TypeGuard["DebugLinkEraseSdCard"]:
             return isinstance(msg, cls)
 
-    class DebugLinkWatchLayout(protobuf.MessageType):
-        watch: "bool | None"
-
-        def __init__(
-            self,
-            *,
-            watch: "bool | None" = None,
-        ) -> None:
-            pass
-
-        @classmethod
-        def is_type_of(cls, msg: Any) -> TypeGuard["DebugLinkWatchLayout"]:
-            return isinstance(msg, cls)
-
-    class DebugLinkResetDebugEvents(protobuf.MessageType):
-
-        @classmethod
-        def is_type_of(cls, msg: Any) -> TypeGuard["DebugLinkResetDebugEvents"]:
-            return isinstance(msg, cls)
-
     class EosGetPublicKey(protobuf.MessageType):
         address_n: "list[int]"
         show_display: "bool | None"

core/src/trezor/pin.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from typing import Any
 
-    from trezor.ui.layouts.common import ProgressLayout
+    from trezor.ui import ProgressLayout
 
 _previous_seconds: int | None = None
 _previous_remaining: str | None = None

core/src/trezor/ui/__init__.py

@@ -1,9 +1,21 @@
 # pylint: disable=wrong-import-position
 import utime
+from micropython import const
 from trezorui import Display
-from typing import TYPE_CHECKING, Any, Awaitable, Generator
+from typing import TYPE_CHECKING
 
-from trezor import loop, utils
+from trezor import io, loop, utils, workflow
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Generator, Generic, Iterator, TypeVar
+
+    from trezorui2 import LayoutObj, UiResult
+
+    T = TypeVar("T", covariant=True)
+
+else:
+    T = 0
+    Generic = {T: object}
 
 # all rendering is done through a singleton of `Display`
 display = Display()
@@ -16,8 +28,10 @@ MONO: int = Display.FONT_MONO
 WIDTH: int = Display.WIDTH
 HEIGHT: int = Display.HEIGHT
 
-# channel used to cancel layouts, see `Cancelled` exception
+_REQUEST_ANIMATION_FRAME = const(1)
-layout_chan = loop.chan()
+"""Animation frame timer token.
+See `trezor::ui::layout::base::EventCtx::ANIM_FRAME_TIMER`.
+"""
 
 # allow only one alert at a time to avoid alerts overlapping
 _alert_in_progress = False
@@ -89,219 +103,293 @@ def backlight_fade(val: int, delay: int = 14000, step: int = 15) -> None:
         display.backlight(val)
 
 
-# Component events.  Should be different from `io.TOUCH_*` events.
+class Shutdown(Exception):
-# Event dispatched when components should draw to the display, if they are
+    pass
-# marked for re-paint.
-RENDER = const(-255)
-# Event dispatched when components should mark themselves for re-painting.
-REPAINT = const(-256)
-
-# How long, in milliseconds, should the layout rendering task sleep between
-# the render calls.
-_RENDER_DELAY_MS = const(10)
 
 
-class Component:
+SHUTDOWN = Shutdown()
+
+CURRENT_LAYOUT: "Layout | ProgressLayout | None" = None
+
+
+def set_current_layout(layout: "Layout | ProgressLayout | None") -> None:
+    """Set the current global layout.
+
+    All manipulation of the global `CURRENT_LAYOUT` MUST go through this function.
+    It ensures that the transitions are always to/from None (so that there are never
+    two layouts in RUNNING state), and that the debug UI is notified of the change.
     """
-    Abstract class.
+    global CURRENT_LAYOUT
 
-    Components are GUI classes that inherit `Component` and form a tree, with a
+    # all transitions must be to/from None
-    `Layout` at the root, and other components underneath.  Components that
+    assert (CURRENT_LAYOUT is None) == (layout is not None)
-    have children, and therefore need to dispatch events to them, usually
+
-    override the `dispatch` method.  Leaf components usually override the event
+    CURRENT_LAYOUT = layout
-    methods (`on_*`).  Components signal a completion to the layout by raising
+
-    an instance of `Result`.
+    if __debug__ and not isinstance(layout, ProgressLayout):
+        from apps.debug import notify_layout_change
+
+        notify_layout_change(layout)
+
+
+class Layout(Generic[T]):
+    """Python-side handler and runner for the Rust based layouts.
+
+    Wrap a `LayoutObj` instance in `Layout` to be able to display the layout, run its
+    event loop, and take part in global layout management. See
+    [docs/core/misc/layout-lifecycle.md] for details.
     """
 
-    def __init__(self) -> None:
+    BACKLIGHT_LEVEL = style.BACKLIGHT_NORMAL
-        self.repaint = True
-
-    def dispatch(self, event: int, x: int, y: int) -> None:
-        if event is RENDER:
-            self.on_render()
-        elif utils.USE_BUTTON and event is io.BUTTON_PRESSED:
-            self.on_button_pressed(x)
-        elif utils.USE_BUTTON and event is io.BUTTON_RELEASED:
-            self.on_button_released(x)
-        elif utils.USE_TOUCH and event is io.TOUCH_START:
-            self.on_touch_start(x, y)
-        elif utils.USE_TOUCH and event is io.TOUCH_MOVE:
-            self.on_touch_move(x, y)
-        elif utils.USE_TOUCH and event is io.TOUCH_END:
-            self.on_touch_end(x, y)
-        elif event is REPAINT:
-            self.repaint = True
-
-    def on_touch_start(self, x: int, y: int) -> None:
-        pass
-
-    def on_touch_move(self, x: int, y: int) -> None:
-        pass
-
-    def on_touch_end(self, x: int, y: int) -> None:
-        pass
-
-    def on_button_pressed(self, button_number: int) -> None:
-        pass
-
-    def on_button_released(self, button_number: int) -> None:
-        pass
-
-    def on_render(self) -> None:
-        pass
 
     if __debug__:
 
-        def read_content_into(self, content_store: list[str]) -> None:
+        @staticmethod
-            content_store.clear()
+        def _trace(layout: LayoutObj) -> str:
-            content_store.append(self.__class__.__name__)
+            tokens = []
 
+            def callback(*args: str) -> None:
+                tokens.extend(args)
 
-class Result(Exception):
+            layout.trace(callback)
-    """
+            return "".join(tokens)
-    When components want to trigger layout completion, they do so through
-    raising an instance of `Result`.
 
-    See `Layout.__iter__` for details.
+        def __str__(self) -> str:
-    """
+            return f"{repr(self)}({self._trace(self.layout)[:150]})"
 
-    def __init__(self, value: Any) -> None:
+    def __init__(self, layout: LayoutObj[T]) -> None:
-        super().__init__()
+        """Set up a layout."""
-        self.value = value
+        self.layout = layout
+        self.tasks: set[loop.Task] = set()
+        self.timers: dict[int, loop.Task] = {}
+        self.result_box = loop.mailbox()
 
+    def is_ready(self) -> bool:
+        """True if the layout is in READY state."""
+        return CURRENT_LAYOUT is not self and self.result_box.is_empty()
 
-class Cancelled(Exception):
+    def is_running(self) -> bool:
-    """
+        """True if the layout is in RUNNING state."""
-    Layouts can be explicitly cancelled.  This usually happens when another
+        return CURRENT_LAYOUT is self
-    layout starts, because only one layout can be running at the same time,
-    and is done by raising `Cancelled` on the cancelled layout.  Layouts
-    should always re-raise such exceptions.
 
-    See `Layout.__iter__` for details.
+    def is_stopped(self) -> bool:
-    """
+        """True if the layout is in STOPPED state."""
+        return CURRENT_LAYOUT is not self and not self.result_box.is_empty()
 
+    def start(self) -> None:
+        """Start the layout, stopping any other RUNNING layout.
 
-class Layout:
+        If the layout is already RUNNING, do nothing. If the layout is STOPPED, fail.
-    """
-    Abstract class.
-
-    Layouts are top-level components.  Only one layout can be running at the
-    same time.  Layouts provide asynchronous interface, so a running task can
-    wait for the layout to complete.  Layouts complete when a `Result` is
-    raised, usually from some of the child components.
-    """
-
-    async def __iter__(self) -> Any:
         """
-        Run the layout and wait until it completes.  Returns the result value.
+        global CURRENT_LAYOUT
-        Usually not overridden.
-        """
-        if __debug__:
-            # we want to call notify_layout_change() when the rendering is done;
-            # but only the first time the layout is awaited. Here we indicate that we
-            # are being awaited, and in handle_rendering() we send the appropriate event
-            self.should_notify_layout_change = True
 
-        value = None
+        # do nothing if we are already running
+        if self.is_running():
+            return
+
+        # make sure we are not restarted before picking the previous result
+        assert self.is_ready()
+
+        # set up the global layout, shutting down any competitors
+        # (caller should still call `workflow.close_others()` to ensure that someone
+        # else will not just shut us down immediately)
+        if CURRENT_LAYOUT is not None:
+            CURRENT_LAYOUT.stop()
+
+        assert CURRENT_LAYOUT is None
+        set_current_layout(self)
+
+        # attach a timer callback and paint self
+        self.layout.attach_timer_fn(self._set_timer)
+        self._first_paint()
+
+        # spawn all tasks
+        for task in self.create_tasks():
+            self.tasks.add(task)
+            loop.schedule(task)
+
+    def stop(self, _kill_taker: bool = True) -> None:
+        """Stop the layout, moving out of RUNNING state and unsetting self as the
+        current layout.
+
+        The resulting state is either READY (if there is no result to be picked up) or
+        STOPPED.
+
+        When called externally, this kills any tasks that wait for the result, assuming
+        that the external `stop()` is a kill. When called internally, `_kill_taker` is
+        set to False to indicate that a result became available and that the taker
+        should be allowed to pick it up.
+        """
+        global CURRENT_LAYOUT
+
+        # stop all running timers and spawned tasks
+        for timer in self.timers.values():
+            loop.close(timer)
+        for task in self.tasks:
+            if task != loop.this_task:
+                loop.close(task)
+        self.timers.clear()
+        self.tasks.clear()
+
+        # shut down anyone who is waiting for the result
+        if _kill_taker:
+            self.result_box.maybe_close()
+
+        if CURRENT_LAYOUT is self:
+            # fade to black -- backlight is off while no layout is running
+            backlight_fade(style.BACKLIGHT_NONE)
+
+            set_current_layout(None)
+
+    async def get_result(self) -> T:
+        """Wait for, and return, the result of this UI layout."""
+        if self.is_ready():
+            self.start()
+        # else we are (a) still running or (b) already stopped
         try:
-            # If any other layout is running (waiting on the layout channel),
+            return await self.result_box
-            # we close it with the Cancelled exception, and wait until it is
+        finally:
-            # closed, just to be sure.
+            self.stop()
-            if layout_chan.takers:
-                await layout_chan.put(Cancelled())
-            # Now, no other layout should be running.  In a loop, we create new
-            # layout tasks and execute them in parallel, while waiting on the
-            # layout channel.  This allows other layouts to cancel us, and the
-            # layout tasks to trigger restart by exiting (new tasks are created
-            # and we continue, because we are in a loop).
-            while True:
-                await loop.race(layout_chan.take(), *self.create_tasks())
-        except Result as result:
-            # Result exception was raised, this means this layout is complete.
-            value = result.value
-        return value
 
-    if TYPE_CHECKING:
+    def request_complete_repaint(self) -> None:
+        """Request a complete repaint of the layout."""
+        msg = self.layout.request_complete_repaint()
+        assert msg is None
 
-        def __await__(self) -> Generator:
+    def _paint(self) -> None:
-            return self.__iter__()  # type: ignore [Expression of type "Coroutine[Any, Any, Any]" cannot be assigned to return type "Generator[Unknown, Unknown, Unknown]"]
+        """Paint the layout and ensure that homescreen cache is properly invalidated."""
+        import storage.cache as storage_cache
 
-    else:
+        painted = self.layout.paint()
-        __await__ = __iter__
-
-    def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
-        """
-        Called from `__iter__`.  Creates and returns a sequence of tasks that
-        run this layout.  Tasks are executed in parallel.  When one of them
-        returns, the others are closed and `create_tasks` is called again.
-
-        Usually overridden to add another tasks to the list."""
-        tasks = (self.handle_rendering(),)
-        if utils.USE_BUTTON:
-            tasks = tasks + (self.handle_button(),)
-        if utils.USE_TOUCH:
-            tasks = tasks + (self.handle_touch(),)
-        return tasks
-
-    def handle_touch(self) -> Generator:
-        """Task that is waiting for the user input."""
-        touch = loop.wait(io.TOUCH)
-        while True:
-            # Using `yield` instead of `await` to avoid allocations.
-            event, x, y = yield touch
-            workflow.idle_timer.touch()
-            self.dispatch(event, x, y)
-            # We dispatch a render event right after the touch.  Quick and dirty
-            # way to get the lowest input-to-render latency.
-            self.dispatch(RENDER, 0, 0)
-
-    def handle_button(self) -> Generator:
-        """Task that is waiting for the user input."""
-        button = loop.wait(io.BUTTON)
-        while True:
-            event, button_num = yield button
-            workflow.idle_timer.touch()
-            self.dispatch(event, button_num, 0)
-            self.dispatch(RENDER, 0, 0)
-
-    def _before_render(self) -> None:
-        # Before the first render, we dim the display.
-        backlight_fade(style.BACKLIGHT_NONE)
-        # Clear the screen of any leftovers, make sure everything is marked for
-        # repaint (we can be running the same layout instance multiple times)
-        # and paint it.
-        display.clear()
-        self.dispatch(REPAINT, 0, 0)
-        self.dispatch(RENDER, 0, 0)
-
-        if __debug__ and self.should_notify_layout_change:
-            from apps.debug import notify_layout_change
-
-            # notify about change and do not notify again until next await.
-            # (handle_rendering might be called multiple times in a single await,
-            # because of the endless loop in __iter__)
-            self.should_notify_layout_change = False
-            notify_layout_change(self)
-
-        # Display is usually refreshed after every loop step, but here we are
-        # rendering everything synchronously, so refresh it manually and turn
-        # the brightness on again.
         refresh()
+        if storage_cache.homescreen_shown is not None and painted:
+            storage_cache.homescreen_shown = None
+
+    def _first_paint(self) -> None:
+        """Paint the layout for the first time after starting it.
+
+        This is a separate call in order for homescreens to be able to override and not
+        paint when the screen contents are still valid.
+        """
+        # Clear the screen of any leftovers.
+        self.request_complete_repaint()
+        self._paint()
+
+        # Turn the brightness on.
         backlight_fade(self.BACKLIGHT_LEVEL)
 
-    def handle_rendering(self) -> loop.Task:  # type: ignore [awaitable-is-generator]
+    def _set_timer(self, token: int, deadline: int) -> None:
-        """Task that is rendering the layout in a busy loop."""
+        """Timer callback for Rust layouts."""
-        self._before_render()
+
-        sleep = self.RENDER_SLEEP
+        async def timer_task() -> None:
-        while True:
+            self.timers.pop(token)
-            # Wait for a couple of ms and render the layout again.  Because
+            result = self.layout.timer(token)
-            # components use re-paint marking, they do not really draw on the
+            self._paint()
-            # display needlessly.  Using `yield` instead of `await` to avoid allocations.
+            if result is not None:
-            # TODO: remove the busy loop
+                self.result_box.put(result)
-            yield sleep
+
-            self.dispatch(RENDER, 0, 0)
+        if token == _REQUEST_ANIMATION_FRAME and token in self.timers:
+            # do not schedule another animation frame if one is already scheduled
+            return
+
+        assert token not in self.timers
+        task = timer_task()
+        self.timers[token] = task
+        loop.schedule(task, token, deadline)
+
+    def _emit_message(self, msg: Any) -> None:
+        """Process a message coming out of the Rust layout. Set is as a result and shut
+        down the layout if appropriate, do nothing otherwise."""
+        if msg is None:
+            return
+
+        # when emitting a message, there should not be another one already waiting
+        assert self.result_box.is_empty()
+
+        self.stop(_kill_taker=False)
+
+        self.result_box.put(msg)
+        raise SHUTDOWN
+
+    def create_tasks(self) -> Iterator[loop.Task]:
+        """Set up background tasks for a layout.
+
+        Called from `start()`. Creates and yields a list of background tasks, typically
+        event handlers for different interfaces.
+
+        Override and then `yield from super().create_tasks()` to add more tasks."""
+        if utils.USE_BUTTON:
+            yield self._handle_input_iface(io.BUTTON, self.layout.button_event)
+        if utils.USE_TOUCH:
+            yield self._handle_input_iface(io.TOUCH, self.layout.touch_event)
+
+    def _handle_input_iface(
+        self, iface: int, event_call: Callable[..., object]
+    ) -> Generator:
+        """Task that is waiting for the user input."""
+        touch = loop.wait(iface)
+        try:
+            while True:
+                # Using `yield` instead of `await` to avoid allocations.
+                event = yield touch
+                workflow.idle_timer.touch()
+                msg = event_call(*event)
+                self._emit_message(msg)
+                self.layout.paint()
+        except Shutdown:
+            return
+        finally:
+            touch.close()
 
 
-def wait_until_layout_is_running() -> Awaitable[None]:  # type: ignore [awaitable-is-generator]
+class ProgressLayout:
-    while not layout_chan.takers:
+    """Progress layout.
-        yield
+
+    Simplified version of the general Layout object, for the purpose of showing spinners
+    and loaders that are shown "in the background" of a running workflow. Does not run
+    background tasks, does not respond to timers.
+
+    Participates in global layout management. This is to track whether the progress bar
+    is currently displayed, who needs to redraw and when.
+    """
+
+    def __init__(self, layout: LayoutObj[UiResult]) -> None:
+        self.layout = layout
+
+    def report(self, value: int, description: str | None = None) -> None:
+        """Report a progress step.
+
+        Starts the layout if it is not running.
+
+        `value` can be in range from 0 to 1000.
+        """
+        if CURRENT_LAYOUT is not self:
+            self.start()
+
+        if utils.DISABLE_ANIMATION:
+            return
+
+        msg = self.layout.progress_event(value, description or "")
+        assert msg is None
+        self.layout.paint()
+        refresh()
+
+    def start(self) -> None:
+        global CURRENT_LAYOUT
+
+        if CURRENT_LAYOUT is not self and CURRENT_LAYOUT is not None:
+            CURRENT_LAYOUT.stop()
+
+        assert CURRENT_LAYOUT is None
+        CURRENT_LAYOUT = self
+
+        self.layout.request_complete_repaint()
+        self.layout.paint()
+        backlight_fade(style.BACKLIGHT_NORMAL)
+        refresh()
+
+    def stop(self) -> None:
+        global CURRENT_LAYOUT
+
+        if CURRENT_LAYOUT is self:
+            CURRENT_LAYOUT = None

core/src/trezor/ui/layouts/common.py

@@ -1,41 +1,48 @@
 from typing import TYPE_CHECKING
 
-from trezor import log, workflow
+import trezorui2
+from trezor import log, ui, workflow
 from trezor.enums import ButtonRequestType
 from trezor.messages import ButtonAck, ButtonRequest
-from trezor.wire import context
+from trezor.wire import ActionCancelled, context
 
 if TYPE_CHECKING:
-    from typing import Any, Awaitable, Protocol
+    from typing import TypeVar
 
-    LayoutType = Awaitable[Any]
     PropertyType = tuple[str | None, str | bytes | None]
     ExceptionType = BaseException | type[BaseException]
 
-    class ProgressLayout(Protocol):
+    T = TypeVar("T")
-        def report(self, value: int, description: str | None = None) -> None:
-            ...
 
 
-async def button_request(
+async def _button_request(
     br_type: str,
     code: ButtonRequestType = ButtonRequestType.Other,
-    pages: int | None = None,
+    pages: int = 0,
 ) -> None:
+    workflow.close_others()
     if __debug__:
         log.debug(__name__, "ButtonRequest.type=%s", br_type)
-    workflow.close_others()
+    await context.maybe_call(ButtonRequest(code=code, pages=pages or None), ButtonAck)
-    await context.maybe_call(ButtonRequest(code=code, pages=pages), ButtonAck)
 
 
 async def interact(
-    layout: LayoutType,
+    layout_obj: ui.LayoutObj[T],
-    br_type: str,
+    br_type: str | None,
     br_code: ButtonRequestType = ButtonRequestType.Other,
-) -> Any:
+    raise_on_cancel: ExceptionType | None = ActionCancelled,
-    pages = None
+) -> T:
-    if hasattr(layout, "page_count") and layout.page_count() > 1:  # type: ignore [Cannot access member "page_count" for type "LayoutType"]
+    # shut down other workflows to prevent them from interfering with the current one
-        # We know for certain how many pages the layout will have
+    workflow.close_others()
-        pages = layout.page_count()  # type: ignore [Cannot access member "page_count" for type "LayoutType"]
+    # start the layout
-    await button_request(br_type, br_code, pages)
+    layout = ui.Layout(layout_obj)
-    return await context.wait(layout)
+    layout.start()
+    # send the button request
+    if br_type is not None:
+        await _button_request(br_type, br_code, layout_obj.page_count())
+    # wait for the layout result
+    result = await context.wait(layout.get_result())
+    # raise an exception if the user cancelled the action
+    if raise_on_cancel is not None and result is trezorui2.CANCELLED:
+        raise raise_on_cancel
+    return result

core/src/trezor/ui/layouts/homescreen.py

@@ -1,6 +1,132 @@
-from trezor import utils
+from typing import TYPE_CHECKING
 
-if utils.UI_LAYOUT == "TT":
+import storage.cache as storage_cache
-    from .tt.homescreen import *  # noqa: F401,F403
+import trezorui2
-elif utils.UI_LAYOUT == "TR":
+from trezor import ui
-    from .tr.homescreen import *  # noqa: F401,F403
+
+if TYPE_CHECKING:
+    from typing import Any, Iterator
+
+    from trezor import loop
+
+
+class HomescreenBase(ui.Layout):
+    RENDER_INDICATOR: object | None = None
+
+    def __init__(self, layout: Any) -> None:
+        super().__init__(layout=layout)
+
+    def _paint(self) -> None:
+        self.layout.paint()
+        ui.refresh()
+
+    def _first_paint(self) -> None:
+        if storage_cache.homescreen_shown is not self.RENDER_INDICATOR:
+            super()._first_paint()
+            storage_cache.homescreen_shown = self.RENDER_INDICATOR
+        # else:
+        #     self._paint()
+
+
+class Homescreen(HomescreenBase):
+    RENDER_INDICATOR = storage_cache.HOMESCREEN_ON
+
+    def __init__(
+        self,
+        label: str | None,
+        notification: str | None,
+        notification_is_error: bool,
+        hold_to_lock: bool,
+    ) -> None:
+        level = 1
+        if notification is not None:
+            notification = notification.rstrip("!")
+            if "COINJOIN" in notification.upper():
+                level = 3
+            elif "EXPERIMENTAL" in notification.upper():
+                level = 2
+            elif notification_is_error:
+                level = 0
+
+        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
+        super().__init__(
+            layout=trezorui2.show_homescreen(
+                label=label,
+                notification=notification,
+                notification_level=level,
+                hold=hold_to_lock,
+                skip_first_paint=skip,
+            ),
+        )
+
+    async def usb_checker_task(self) -> None:
+        from trezor import io, loop
+
+        usbcheck = loop.wait(io.USB_CHECK)
+        while True:
+            is_connected = await usbcheck
+            self.layout.usb_event(is_connected)
+            self.layout.paint()
+            ui.refresh()
+
+    def create_tasks(self) -> Iterator[loop.Task]:
+        yield from super().create_tasks()
+        yield self.usb_checker_task()
+
+
+class Lockscreen(HomescreenBase):
+    RENDER_INDICATOR = storage_cache.LOCKSCREEN_ON
+    BACKLIGHT_LEVEL = ui.style.BACKLIGHT_LOW
+
+    def __init__(
+        self,
+        label: str | None,
+        bootscreen: bool = False,
+        coinjoin_authorized: bool = False,
+    ) -> None:
+        self.bootscreen = bootscreen
+        if bootscreen:
+            self.BACKLIGHT_LEVEL = ui.style.BACKLIGHT_NORMAL
+
+        skip = (
+            not bootscreen and storage_cache.homescreen_shown is self.RENDER_INDICATOR
+        )
+        super().__init__(
+            layout=trezorui2.show_lockscreen(
+                label=label,
+                bootscreen=bootscreen,
+                skip_first_paint=skip,
+                coinjoin_authorized=coinjoin_authorized,
+            ),
+        )
+
+    async def get_result(self) -> Any:
+        result = await super().get_result()
+        if self.bootscreen:
+            self.request_complete_repaint()
+        return result
+
+
+class Busyscreen(HomescreenBase):
+    RENDER_INDICATOR = storage_cache.BUSYSCREEN_ON
+
+    def __init__(self, delay_ms: int) -> None:
+        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
+        super().__init__(
+            layout=trezorui2.show_progress_coinjoin(
+                title="Waiting for others",
+                indeterminate=True,
+                time_ms=delay_ms,
+                skip_first_paint=skip,
+            )
+        )
+
+    async def get_result(self) -> Any:
+        from apps.base import set_homescreen
+
+        # Handle timeout.
+        result = await super().get_result()
+        assert result == trezorui2.CANCELLED
+        storage_cache.delete(storage_cache.APP_COMMON_BUSY_DEADLINE_MS)
+        set_homescreen()
+        return result

core/src/trezor/ui/layouts/progress.py

@@ -1,6 +1,50 @@
-from trezor import utils
+import trezorui2
+from trezor import ui, utils
 
-if utils.UI_LAYOUT == "TT":
+
-    from .tt.progress import *  # noqa: F401,F403
+def progress(
-elif utils.UI_LAYOUT == "TR":
+    message: str = "PLEASE WAIT",
-    from .tr.progress import *  # noqa: F401,F403
+    description: str | None = None,
+    indeterminate: bool = False,
+) -> ui.ProgressLayout:
+    if utils.MODEL_IS_T2B1 and description is None:
+        description = message + "..."
+        title = ""
+    else:
+        title = message.upper()
+
+    return ui.ProgressLayout(
+        layout=trezorui2.show_progress(
+            title=title,
+            indeterminate=indeterminate,
+            description=description or "",
+        )
+    )
+
+
+def bitcoin_progress(message: str) -> ui.ProgressLayout:
+    return progress(message)
+
+
+def coinjoin_progress(message: str) -> ui.ProgressLayout:
+    return ui.ProgressLayout(
+        layout=trezorui2.show_progress_coinjoin(
+            title=message + "...", indeterminate=False
+        )
+    )
+
+
+def pin_progress(message: str, description: str) -> ui.ProgressLayout:
+    return progress(message, description=description)
+
+
+def monero_keyimage_sync_progress() -> ui.ProgressLayout:
+    return progress("Syncing")
+
+
+def monero_live_refresh_progress() -> ui.ProgressLayout:
+    return progress("Refreshing", indeterminate=True)
+
+
+def monero_transaction_progress_inner() -> ui.ProgressLayout:
+    return progress("Signing transaction")

core/src/trezor/ui/layouts/tr/fido.py

@@ -1,8 +1,8 @@
 import trezorui2
+from trezor import ui
 from trezor.enums import ButtonRequestType
 
 from ..common import interact
-from . import RustLayout
 
 
 async def confirm_fido(
@@ -12,17 +12,13 @@ async def confirm_fido(
     accounts: list[str | None],
 ) -> int:
     """Webauthn confirmation for one or more credentials."""
-    confirm = RustLayout(
+    confirm = trezorui2.confirm_fido(  # type: ignore [Argument missing for parameter "icon_name"]
-        trezorui2.confirm_fido(  # type: ignore [Argument missing for parameter "icon_name"]
+        title=header.upper(),
-            title=header.upper(),
+        app_name=app_name,
-            app_name=app_name,
+        accounts=accounts,
-            accounts=accounts,
-        )
     )
     result = await interact(confirm, "confirm_fido", ButtonRequestType.Other)
 
-    # The Rust side returns either an int or `CANCELLED`. We detect the int situation
-    # and assume cancellation otherwise.
     if isinstance(result, int):
         return result
 
@@ -31,20 +27,15 @@ async def confirm_fido(
     if __debug__ and result is trezorui2.CONFIRMED:
         return 0
 
-    # Late import won't get executed on the happy path.
+    raise RuntimeError  # should not get here, cancellation is handled by `interact`
-    from trezor.wire import ActionCancelled
-
-    raise ActionCancelled
 
 
 async def confirm_fido_reset() -> bool:
-    confirm = RustLayout(
+    confirm = trezorui2.confirm_action(
-        trezorui2.confirm_action(
+        title="FIDO2 RESET",
-            title="FIDO2 RESET",
+        description="Do you really want to erase all credentials?",
-            description="Do you really want to erase all credentials?",
+        action=None,
-            action=None,
+        verb_cancel="",
-            verb_cancel="",
+        verb="CONFIRM",
-            verb="CONFIRM",
-        )
     )
-    return (await confirm) is trezorui2.CONFIRMED
+    return (await ui.Layout(confirm).get_result()) is trezorui2.CONFIRMED

core/src/trezor/ui/layouts/tr/homescreen.py

@@ -1,127 +0,0 @@
-from typing import TYPE_CHECKING
-
-import storage.cache as storage_cache
-import trezorui2
-from trezor import ui
-
-from . import RustLayout
-
-if TYPE_CHECKING:
-    from typing import Any, Tuple
-
-    from trezor import loop
-
-
-class HomescreenBase(RustLayout):
-    RENDER_INDICATOR: object | None = None
-
-    def __init__(self, layout: Any) -> None:
-        super().__init__(layout=layout)
-
-    def _paint(self) -> None:
-        self.layout.paint()
-        ui.refresh()
-
-    def _first_paint(self) -> None:
-        if storage_cache.homescreen_shown is not self.RENDER_INDICATOR:
-            super()._first_paint()
-            storage_cache.homescreen_shown = self.RENDER_INDICATOR
-        else:
-            self._paint()
-
-
-class Homescreen(HomescreenBase):
-    RENDER_INDICATOR = storage_cache.HOMESCREEN_ON
-
-    def __init__(
-        self,
-        label: str | None,
-        notification: str | None,
-        notification_is_error: bool,
-        hold_to_lock: bool,
-    ) -> None:
-        level = 1
-        if notification is not None:
-            notification = notification.rstrip("!")
-            if "EXPERIMENTAL" in notification:
-                level = 2
-            elif notification_is_error:
-                level = 0
-
-        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        super().__init__(
-            layout=trezorui2.show_homescreen(
-                label=label,
-                notification=notification,
-                notification_level=level,
-                hold=hold_to_lock,
-                skip_first_paint=skip,
-            ),
-        )
-
-    async def usb_checker_task(self) -> None:
-        from trezor import io, loop
-
-        usbcheck = loop.wait(io.USB_CHECK)
-        while True:
-            is_connected = await usbcheck
-            self.layout.usb_event(is_connected)
-            self.layout.paint()
-            ui.refresh()
-
-    def create_tasks(self) -> Tuple[loop.AwaitableTask, ...]:
-        return super().create_tasks() + (self.usb_checker_task(),)
-
-
-class Lockscreen(HomescreenBase):
-    RENDER_INDICATOR = storage_cache.LOCKSCREEN_ON
-
-    def __init__(
-        self,
-        label: str | None,
-        bootscreen: bool = False,
-        coinjoin_authorized: bool = False,
-    ) -> None:
-        self.bootscreen = bootscreen
-        skip = (
-            not bootscreen and storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        )
-        super().__init__(
-            layout=trezorui2.show_lockscreen(
-                label=label,
-                bootscreen=bootscreen,
-                skip_first_paint=skip,
-                coinjoin_authorized=coinjoin_authorized,
-            ),
-        )
-
-    async def __iter__(self) -> Any:
-        result = await super().__iter__()
-        if self.bootscreen:
-            self.request_complete_repaint()
-        return result
-
-
-class Busyscreen(HomescreenBase):
-    RENDER_INDICATOR = storage_cache.BUSYSCREEN_ON
-
-    def __init__(self, delay_ms: int) -> None:
-        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        super().__init__(
-            layout=trezorui2.show_progress_coinjoin(
-                title="Waiting for others",
-                indeterminate=True,
-                time_ms=delay_ms,
-                skip_first_paint=skip,
-            )
-        )
-
-    async def __iter__(self) -> Any:
-        from apps.base import set_homescreen
-
-        # Handle timeout.
-        result = await super().__iter__()
-        assert result == trezorui2.CANCELLED
-        storage_cache.delete(storage_cache.APP_COMMON_BUSY_DEADLINE_MS)
-        set_homescreen()
-        return result

core/src/trezor/ui/layouts/tr/progress.py

@@ -1,68 +0,0 @@
-from typing import TYPE_CHECKING
-
-import trezorui2
-from trezor import ui
-
-if TYPE_CHECKING:
-    from typing import Any
-
-    from ..common import ProgressLayout
-
-
-class RustProgress:
-    def __init__(
-        self,
-        layout: Any,
-    ):
-        self.layout = layout
-        self.layout.attach_timer_fn(self.set_timer)
-        self.layout.paint()
-
-    def set_timer(self, token: int, deadline: int) -> None:
-        raise RuntimeError  # progress layouts should not set timers
-
-    def report(self, value: int, description: str | None = None):
-        msg = self.layout.progress_event(value, description or "")
-        assert msg is None
-        self.layout.paint()
-        ui.refresh()
-
-
-def progress(
-    message: str = "PLEASE WAIT",
-    description: str | None = None,
-    indeterminate: bool = False,
-) -> ProgressLayout:
-    return RustProgress(
-        layout=trezorui2.show_progress(
-            title=message.upper(),
-            indeterminate=indeterminate,
-            description=description or "",
-        )
-    )
-
-
-def bitcoin_progress(description: str) -> ProgressLayout:
-    return progress("", description)
-
-
-def coinjoin_progress(message: str) -> ProgressLayout:
-    return RustProgress(
-        layout=trezorui2.show_progress_coinjoin(title=message, indeterminate=False)
-    )
-
-
-def pin_progress(message: str, description: str) -> ProgressLayout:
-    return progress(message, description)
-
-
-def monero_keyimage_sync_progress() -> ProgressLayout:
-    return progress("", "Syncing...")
-
-
-def monero_live_refresh_progress() -> ProgressLayout:
-    return progress("", "Refreshing...", indeterminate=True)
-
-
-def monero_transaction_progress_inner() -> ProgressLayout:
-    return progress("", "Signing transaction...")

core/src/trezor/ui/layouts/tr/recovery.py

@@ -1,33 +1,38 @@
-from typing import Callable, Iterable
+from typing import Awaitable, Callable, Iterable
 
 import trezorui2
+from trezor import ui
 from trezor.enums import ButtonRequestType
 
 from ..common import interact
-from . import RustLayout, raise_if_not_confirmed, show_warning
+from . import show_warning
 
 
 async def request_word_count(dry_run: bool) -> int:
     count = await interact(
-        RustLayout(trezorui2.select_word_count(dry_run=dry_run)),
+        trezorui2.select_word_count(dry_run=dry_run),
-        "word_count",
+        "recovery_word_count",
         ButtonRequestType.MnemonicWordCount,
     )
     # It can be returning a string (for example for __debug__ in tests)
     return int(count)
 
 
-async def request_word(word_index: int, word_count: int, is_slip39: bool) -> str:
+async def request_word(
-    from trezor.wire.context import wait
+    word_index: int, word_count: int, is_slip39: bool, send_button_request: bool
-
+) -> str:
     prompt = f"WORD {word_index + 1} OF {word_count}"
 
     if is_slip39:
-        word_choice = RustLayout(trezorui2.request_slip39(prompt=prompt))
+        keyboard = trezorui2.request_slip39(prompt=prompt)
     else:
-        word_choice = RustLayout(trezorui2.request_bip39(prompt=prompt))
+        keyboard = trezorui2.request_bip39(prompt=prompt)
 
-    word: str = await wait(word_choice)
+    word: str = await interact(
+        keyboard,
+        "mnemonic" if send_button_request else None,
+        ButtonRequestType.MnemonicInput,
+    )
     return word
 
 
@@ -39,22 +44,20 @@ async def show_remaining_shares(
     raise NotImplementedError
 
 
-async def show_group_share_success(share_index: int, group_index: int) -> None:
+def show_group_share_success(
-    await raise_if_not_confirmed(
+    share_index: int, group_index: int
-        interact(
+) -> Awaitable[ui.UiResult]:
-            RustLayout(
+    return interact(
-                trezorui2.show_group_share_success(
+        trezorui2.show_group_share_success(
-                    lines=[
+            lines=[
-                        "You have entered",
+                "You have entered",
-                        f"Share {share_index + 1}",
+                f"Share {share_index + 1}",
-                        "from",
+                "from",
-                        f"Group {group_index + 1}",
+                f"Group {group_index + 1}",
-                    ],
+            ],
-                )
+        ),
-            ),
+        "share_success",
-            "share_success",
+        ButtonRequestType.Other,
-            ButtonRequestType.Other,
-        )
     )
 
 
@@ -77,29 +80,28 @@ async def continue_recovery(
     if subtext:
         text += f"\n\n{subtext}"
 
-    homepage = RustLayout(
+    homepage = trezorui2.confirm_recovery(
-        trezorui2.confirm_recovery(
+        title="",
-            title="",
+        description=text,
-            description=text,
+        button=button_label.upper(),
-            button=button_label.upper(),
+        info_button=False,
-            info_button=False,
+        dry_run=dry_run,
-            dry_run=dry_run,
+        show_info=show_info,  # type: ignore [No parameter named "show_info"]
-            show_info=show_info,  # type: ignore [No parameter named "show_info"]
-        )
     )
     result = await interact(
         homepage,
         "recovery",
         ButtonRequestType.RecoveryHomepage,
+        raise_on_cancel=None,
     )
     return result is trezorui2.CONFIRMED
 
 
-async def show_recovery_warning(
+def show_recovery_warning(
     br_type: str,
     content: str,
     subheader: str | None = None,
     button: str = "TRY AGAIN",
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[ui.UiResult]:
-    await show_warning(br_type, content, subheader, button, br_code)
+    return show_warning(br_type, content, subheader, button, br_code)

core/src/trezor/ui/layouts/tr/reset.py

@@ -2,15 +2,14 @@ from typing import TYPE_CHECKING
 
 import trezorui2
 from trezor.enums import ButtonRequestType
-from trezor.wire import ActionCancelled
 
 from ..common import interact
-from . import RustLayout, confirm_action, show_warning
+from . import confirm_action, show_warning
 
 CONFIRMED = trezorui2.CONFIRMED  # global_import_cache
 
 if TYPE_CHECKING:
-    from typing import Sequence
+    from typing import Awaitable, Sequence
 
     from trezor.enums import BackupType
 
@@ -47,13 +46,12 @@ async def show_share_words(
         )
 
         result = await interact(
-            RustLayout(
+            trezorui2.show_share_words(  # type: ignore [Arguments missing for parameters]
-                trezorui2.show_share_words(  # type: ignore [Arguments missing for parameters]
+                share_words=share_words,  # type: ignore [No parameter named "share_words"]
-                    share_words=share_words,  # type: ignore [No parameter named "share_words"]
-                )
             ),
             br_type,
             br_code,
+            raise_on_cancel=None,
         )
         if result is CONFIRMED:
             break
@@ -76,7 +74,6 @@ async def select_word(
     group_index: int | None = None,
 ) -> str:
     from trezor.strings import format_ordinal
-    from trezor.wire.context import wait
 
     # It may happen (with a very low probability)
     # that there will be less than three unique words to choose from.
@@ -85,14 +82,13 @@ async def select_word(
     while len(words) < 3:
         words.append(words[-1])
 
-    result = await wait(
+    result = await interact(
-        RustLayout(
+        trezorui2.select_word(
-            trezorui2.select_word(
+            title="",
-                title="",
+            description=f"SELECT {format_ordinal(checked_index + 1).upper()} WORD",
-                description=f"SELECT {format_ordinal(checked_index + 1).upper()} WORD",
+            words=(words[0].lower(), words[1].lower(), words[2].lower()),
-                words=(words[0].lower(), words[1].lower(), words[2].lower()),
+        ),
-            )
+        None,
-        )
     )
     if __debug__ and isinstance(result, str):
         return result
@@ -119,20 +115,16 @@ async def slip39_show_checklist(step: int, backup_type: BackupType) -> None:
         )
     )
 
-    result = await interact(
+    await interact(
-        RustLayout(
+        trezorui2.show_checklist(
-            trezorui2.show_checklist(
+            title="BACKUP CHECKLIST",
-                title="BACKUP CHECKLIST",
+            button="CONTINUE",
-                button="CONTINUE",
+            active=step,
-                active=step,
+            items=items,
-                items=items,
-            )
         ),
         "slip39_checklist",
         ButtonRequestType.ResetDevice,
     )
-    if result is not CONFIRMED:
-        raise ActionCancelled
 
 
 async def _prompt_number(
@@ -142,13 +134,11 @@ async def _prompt_number(
     max_count: int,
     br_name: str,
 ) -> int:
-    num_input = RustLayout(
+    num_input = trezorui2.request_number(
-        trezorui2.request_number(
+        title=title.upper(),
-            title=title.upper(),
+        count=count,
-            count=count,
+        min_count=min_count,
-            min_count=min_count,
+        max_count=max_count,
-            max_count=max_count,
-        )
     )
 
     result = await interact(
@@ -156,8 +146,15 @@ async def _prompt_number(
         br_name,
         ButtonRequestType.ResetDevice,
     )
+    if __debug__:
+        if not isinstance(result, tuple):
+            # handle a debuglink confirmation. According to comments on TT side,
+            # debuglink is not sending the value, just a confirmation, and never
+            # modifies the initial count, so let's use that.
+            result = result, count
 
-    return int(result)
+    _status, value = result
+    return value
 
 
 async def slip39_prompt_threshold(
@@ -218,12 +215,12 @@ async def slip39_prompt_number_of_shares(group_id: int | None = None) -> int:
     )
 
 
-async def slip39_advanced_prompt_number_of_groups() -> int:
+def slip39_advanced_prompt_number_of_groups() -> Awaitable[int]:
     count = 5
     min_count = 2
     max_count = 16
 
-    return await _prompt_number(
+    return _prompt_number(
         "NUMBER OF GROUPS",
         count,
         min_count,
@@ -232,12 +229,12 @@ async def slip39_advanced_prompt_number_of_groups() -> int:
     )
 
 
-async def slip39_advanced_prompt_group_threshold(num_of_groups: int) -> int:
+def slip39_advanced_prompt_group_threshold(num_of_groups: int) -> Awaitable[int]:
     count = num_of_groups // 2 + 1
     min_count = 1
     max_count = num_of_groups
 
-    return await _prompt_number(
+    return _prompt_number(
         "GROUP THRESHOLD",
         count,
         min_count,
@@ -246,8 +243,8 @@ async def slip39_advanced_prompt_group_threshold(num_of_groups: int) -> int:
     )
 
 
-async def show_warning_backup(slip39: bool) -> None:
+def show_warning_backup(slip39: bool) -> Awaitable[trezorui2.UiResult]:
-    await show_warning(
+    return show_warning(
         "backup_warning",
         "REMEMBER",
         "Never make a digital copy of your backup or upload it online!",
@@ -256,8 +253,8 @@ async def show_warning_backup(slip39: bool) -> None:
     )
 
 
-async def show_success_backup() -> None:
+def show_success_backup() -> Awaitable[trezorui2.UiResult]:
-    await confirm_action(
+    return confirm_action(
         "success_backup",
         "BACKUP IS DONE",
         description="Keep it safe!",
@@ -267,14 +264,14 @@ async def show_success_backup() -> None:
     )
 
 
-async def show_reset_warning(
+def show_reset_warning(
     br_type: str,
     content: str,
     subheader: str | None = None,
     button: str = "TRY AGAIN",
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[trezorui2.UiResult]:
-    await show_warning(
+    return show_warning(
         br_type,
         subheader or "",
         content,

core/src/trezor/ui/layouts/tt/fido.py

@@ -1,56 +1,8 @@
-from typing import TYPE_CHECKING
-
 import trezorui2
+from trezor import ui
 from trezor.enums import ButtonRequestType
 
 from ..common import interact
-from . import RustLayout
-
-if TYPE_CHECKING:
-    from trezor.loop import AwaitableTask
-
-
-if __debug__:
-    from trezor import io, ui
-
-    from ... import Result
-
-    class _RustFidoLayoutImpl(RustLayout):
-        def create_tasks(self) -> tuple[AwaitableTask, ...]:
-            from trezor import utils
-
-            tasks = (
-                self.handle_timers(),
-                self.handle_swipe(),
-                self.handle_debug_confirm(),
-            )
-            if utils.USE_TOUCH:
-                tasks = tasks + (self.handle_touch(),)
-            if utils.USE_BUTTON:
-                tasks = tasks + (self.handle_button(),)
-            return tasks
-
-        async def handle_debug_confirm(self) -> None:
-            from apps.debug import result_signal
-
-            _event_id, result = await result_signal()
-            if result is not trezorui2.CONFIRMED:
-                raise Result(result)
-
-            for event, x, y in (
-                (io.TOUCH_START, 220, 220),
-                (io.TOUCH_END, 220, 220),
-            ):
-                msg = self.layout.touch_event(event, x, y)
-                self.layout.paint()
-                ui.refresh()
-                if msg is not None:
-                    raise Result(msg)
-
-    _RustFidoLayout = _RustFidoLayoutImpl
-
-else:
-    _RustFidoLayout = RustLayout
 
 
 async def confirm_fido(
@@ -60,16 +12,30 @@ async def confirm_fido(
     accounts: list[str | None],
 ) -> int:
     """Webauthn confirmation for one or more credentials."""
-    confirm = _RustFidoLayout(
+    confirm = trezorui2.confirm_fido(
-        trezorui2.confirm_fido(
+        title=header.upper(),
-            title=header.upper(),
+        app_name=app_name,
-            app_name=app_name,
+        icon_name=icon_name,
-            icon_name=icon_name,
+        accounts=accounts,
-            accounts=accounts,
-        )
     )
     result = await interact(confirm, "confirm_fido", ButtonRequestType.Other)
 
+    if __debug__ and result is trezorui2.CONFIRMED:
+        # debuglink will directly inject a CONFIRMED message which we need to handle
+        # by playing back a click to the Rust layout and getting out the selected number
+        # that way
+        from trezor import io
+
+        msg = confirm.touch_event(io.TOUCH_START, 220, 220)
+        assert msg is None
+        confirm.paint()
+        ui.refresh()
+        msg = confirm.touch_event(io.TOUCH_END, 220, 220)
+        confirm.paint()
+        ui.refresh()
+        assert isinstance(msg, int)
+        return msg
+
     # The Rust side returns either an int or `CANCELLED`. We detect the int situation
     # and assume cancellation otherwise.
     if isinstance(result, int):
@@ -82,7 +48,7 @@ async def confirm_fido(
 
 
 async def confirm_fido_reset() -> bool:
-    confirm = RustLayout(
+    confirm = ui.Layout(
         trezorui2.confirm_action(
             title="FIDO2 RESET",
             action="erase all credentials?",
@@ -90,4 +56,4 @@ async def confirm_fido_reset() -> bool:
             reverse=True,
         )
     )
-    return (await confirm) is trezorui2.CONFIRMED
+    return (await confirm.get_result()) is trezorui2.CONFIRMED

core/src/trezor/ui/layouts/tt/homescreen.py

@@ -1,148 +0,0 @@
-from typing import TYPE_CHECKING
-
-import storage.cache as storage_cache
-from trezor import ui, utils
-
-import trezorui2
-from trezor import ui
-
-from . import RustLayout
-
-if TYPE_CHECKING:
-    from typing import Any, Tuple
-
-    from trezor import loop
-
-
-class HomescreenBase(RustLayout):
-    RENDER_INDICATOR: object | None = None
-
-    def __init__(self, layout: Any) -> None:
-        super().__init__(layout=layout)
-
-    def _paint(self) -> None:
-        self.layout.paint()
-        ui.refresh()
-
-    def _first_paint(self) -> None:
-        if storage_cache.homescreen_shown is not self.RENDER_INDICATOR:
-            super()._first_paint()
-            storage_cache.homescreen_shown = self.RENDER_INDICATOR
-        else:
-            self._paint()
-
-    if __debug__:
-        # In __debug__ mode, ignore {confirm,swipe,input}_signal.
-        def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
-            tasks = (
-                self.handle_timers(),
-                self.handle_click_signal(),  # so we can receive debug events
-            )
-            if utils.USE_TOUCH:
-                tasks = tasks + (self.handle_touch(),)
-            if utils.USE_BUTTON:
-                tasks = tasks + (self.handle_button(),)
-            return tasks
-
-
-class Homescreen(HomescreenBase):
-    RENDER_INDICATOR = storage_cache.HOMESCREEN_ON
-
-    def __init__(
-        self,
-        label: str | None,
-        notification: str | None,
-        notification_is_error: bool,
-        hold_to_lock: bool,
-    ) -> None:
-        level = 1
-        if notification is not None:
-            notification = notification.rstrip("!")
-            if "COINJOIN" in notification.upper():
-                level = 3
-            elif "EXPERIMENTAL" in notification.upper():
-                level = 2
-            elif notification_is_error:
-                level = 0
-
-        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        super().__init__(
-            layout=trezorui2.show_homescreen(
-                label=label,
-                notification=notification,
-                notification_level=level,
-                hold=hold_to_lock,
-                skip_first_paint=skip,
-            ),
-        )
-
-    async def usb_checker_task(self) -> None:
-        from trezor import io, loop
-
-        usbcheck = loop.wait(io.USB_CHECK)
-        while True:
-            is_connected = await usbcheck
-            self.layout.usb_event(is_connected)
-            self.layout.paint()
-            ui.refresh()
-
-    def create_tasks(self) -> Tuple[loop.AwaitableTask, ...]:
-        return super().create_tasks() + (self.usb_checker_task(),)
-
-
-class Lockscreen(HomescreenBase):
-    RENDER_INDICATOR = storage_cache.LOCKSCREEN_ON
-    BACKLIGHT_LEVEL = ui.style.BACKLIGHT_LOW
-
-    def __init__(
-        self,
-        label: str | None,
-        bootscreen: bool = False,
-        coinjoin_authorized: bool = False,
-    ) -> None:
-        self.bootscreen = bootscreen
-        if bootscreen:
-            self.BACKLIGHT_LEVEL = ui.style.BACKLIGHT_NORMAL
-
-        skip = (
-            not bootscreen and storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        )
-        super().__init__(
-            layout=trezorui2.show_lockscreen(
-                label=label,
-                bootscreen=bootscreen,
-                skip_first_paint=skip,
-                coinjoin_authorized=coinjoin_authorized,
-            ),
-        )
-
-    async def __iter__(self) -> Any:
-        result = await super().__iter__()
-        if self.bootscreen:
-            self.request_complete_repaint()
-        return result
-
-
-class Busyscreen(HomescreenBase):
-    RENDER_INDICATOR = storage_cache.BUSYSCREEN_ON
-
-    def __init__(self, delay_ms: int) -> None:
-        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        super().__init__(
-            layout=trezorui2.show_progress_coinjoin(
-                title="Waiting for others",
-                indeterminate=True,
-                time_ms=delay_ms,
-                skip_first_paint=skip,
-            )
-        )
-
-    async def __iter__(self) -> Any:
-        from apps.base import set_homescreen
-
-        # Handle timeout.
-        result = await super().__iter__()
-        assert result == trezorui2.CANCELLED
-        storage_cache.delete(storage_cache.APP_COMMON_BUSY_DEADLINE_MS)
-        set_homescreen()
-        return result

core/src/trezor/ui/layouts/tt/progress.py

@@ -1,70 +0,0 @@
-from typing import TYPE_CHECKING
-
-import trezorui2
-from trezor import ui
-
-if TYPE_CHECKING:
-    from typing import Any
-
-    from ..common import ProgressLayout
-
-
-class RustProgress:
-    def __init__(
-        self,
-        layout: Any,
-    ):
-        self.layout = layout
-        ui.backlight_fade(ui.style.BACKLIGHT_DIM)
-        self.layout.attach_timer_fn(self.set_timer)
-        self.layout.paint()
-        ui.backlight_fade(ui.style.BACKLIGHT_NORMAL)
-
-    def set_timer(self, token: int, deadline: int) -> None:
-        raise RuntimeError  # progress layouts should not set timers
-
-    def report(self, value: int, description: str | None = None):
-        msg = self.layout.progress_event(value, description or "")
-        assert msg is None
-        self.layout.paint()
-        ui.refresh()
-
-
-def progress(
-    message: str = "PLEASE WAIT",
-    description: str | None = None,
-    indeterminate: bool = False,
-) -> ProgressLayout:
-    return RustProgress(
-        layout=trezorui2.show_progress(
-            title=message.upper(),
-            indeterminate=indeterminate,
-            description=description or "",
-        )
-    )
-
-
-def bitcoin_progress(message: str) -> ProgressLayout:
-    return progress(message)
-
-
-def coinjoin_progress(message: str) -> ProgressLayout:
-    return RustProgress(
-        layout=trezorui2.show_progress_coinjoin(title=message, indeterminate=False)
-    )
-
-
-def pin_progress(message: str, description: str) -> ProgressLayout:
-    return progress(message, description=description)
-
-
-def monero_keyimage_sync_progress() -> ProgressLayout:
-    return progress("SYNCING")
-
-
-def monero_live_refresh_progress() -> ProgressLayout:
-    return progress("REFRESHING", indeterminate=True)
-
-
-def monero_transaction_progress_inner() -> ProgressLayout:
-    return progress("SIGNING TRANSACTION")

core/src/trezor/ui/layouts/tt/recovery.py

@@ -1,22 +1,21 @@
-from typing import Callable, Iterable
+from typing import Awaitable, Callable, Iterable
 
 import trezorui2
+from trezor import ui
 from trezor.enums import ButtonRequestType
-from trezor.wire.context import wait as ctx_wait
 
 from ..common import interact
-from . import RustLayout, raise_if_not_confirmed
 
 CONFIRMED = trezorui2.CONFIRMED  # global_import_cache
 INFO = trezorui2.INFO  # global_import_cache
 
 
 async def _is_confirmed_info(
-    dialog: RustLayout,
+    dialog: ui.LayoutObj,
     info_func: Callable,
 ) -> bool:
     while True:
-        result = await ctx_wait(dialog)
+        result = await interact(dialog, None, raise_on_cancel=None)
 
         if result is trezorui2.INFO:
             await info_func()
@@ -26,27 +25,36 @@ async def _is_confirmed_info(
 
 
 async def request_word_count(dry_run: bool) -> int:
-    selector = RustLayout(trezorui2.select_word_count(dry_run=dry_run))
+    count = await interact(
-    count = await interact(selector, "word_count", ButtonRequestType.MnemonicWordCount)
+        trezorui2.select_word_count(dry_run=dry_run),
+        "word_count",
+        ButtonRequestType.MnemonicWordCount,
+    )
     return int(count)
 
 
-async def request_word(word_index: int, word_count: int, is_slip39: bool) -> str:
+async def request_word(
+    word_index: int, word_count: int, is_slip39: bool, send_button_request: bool
+) -> str:
     prompt = f"Type word {word_index + 1} of {word_count}"
     if is_slip39:
-        keyboard = RustLayout(trezorui2.request_slip39(prompt=prompt))
+        keyboard = trezorui2.request_slip39(prompt=prompt)
     else:
-        keyboard = RustLayout(trezorui2.request_bip39(prompt=prompt))
+        keyboard = trezorui2.request_bip39(prompt=prompt)
 
-    word: str = await ctx_wait(keyboard)
+    word: str = await interact(
+        keyboard,
+        "mnemonic" if send_button_request else None,
+        ButtonRequestType.MnemonicInput,
+    )
     return word
 
 
-async def show_remaining_shares(
+def show_remaining_shares(
     groups: Iterable[tuple[int, tuple[str, ...]]],  # remaining + list 3 words
     shares_remaining: list[int],
     group_threshold: int,
-) -> None:
+) -> Awaitable[trezorui2.UiResult]:
     from trezor import strings
     from trezor.crypto.slip39 import MAX_SHARE_COUNT
 
@@ -68,31 +76,27 @@ async def show_remaining_shares(
             words = "\n".join(group)
             pages.append((title, words))
 
-    await raise_if_not_confirmed(
+    return interact(
-        interact(
+        trezorui2.show_remaining_shares(pages=pages),
-            RustLayout(trezorui2.show_remaining_shares(pages=pages)),
+        "show_shares",
-            "show_shares",
+        ButtonRequestType.Other,
-            ButtonRequestType.Other,
-        )
     )
 
 
-async def show_group_share_success(share_index: int, group_index: int) -> None:
+def show_group_share_success(
-    await raise_if_not_confirmed(
+    share_index: int, group_index: int
-        interact(
+) -> Awaitable[ui.UiResult]:
-            RustLayout(
+    return interact(
-                trezorui2.show_group_share_success(
+        trezorui2.show_group_share_success(
-                    lines=[
+            lines=[
-                        "You have entered",
+                "You have entered",
-                        f"Share {share_index + 1}",
+                f"Share {share_index + 1}",
-                        "from",
+                "from",
-                        f"Group {group_index + 1}",
+                f"Group {group_index + 1}",
-                    ],
+            ],
-                )
+        ),
-            ),
+        "share_success",
-            "share_success",
+        ButtonRequestType.Other,
-            ButtonRequestType.Other,
-        )
     )
 
 
@@ -104,51 +108,50 @@ async def continue_recovery(
     dry_run: bool,
     show_info: bool = False,  # unused on TT
 ) -> bool:
-    from ..common import button_request
-
     if show_info:
         # Show this just one-time
         description = "You'll only have to select the first 2-4 letters of each word."
     else:
         description = subtext or ""
 
-    homepage = RustLayout(
+    homepage = trezorui2.confirm_recovery(
-        trezorui2.confirm_recovery(
+        title=text,
-            title=text,
+        description=description,
-            description=description,
+        button=button_label.upper(),
-            button=button_label.upper(),
+        info_button=info_func is not None,
-            info_button=info_func is not None,
+        dry_run=dry_run,
-            dry_run=dry_run,
-        )
     )
 
-    await button_request("recovery", ButtonRequestType.RecoveryHomepage)
+    send_button_request = True
+    while True:
+        result = await interact(
+            homepage,
+            "recovery" if send_button_request else None,
+            ButtonRequestType.RecoveryHomepage,
+            raise_on_cancel=None,
+        )
 
-    if info_func is not None:
+        if info_func is not None and result is trezorui2.INFO:
-        return await _is_confirmed_info(homepage, info_func)
+            await info_func()
-    else:
+            homepage.request_complete_repaint()
-        result = await ctx_wait(homepage)
+        else:
-        return result is CONFIRMED
+            return result is CONFIRMED
 
 
-async def show_recovery_warning(
+def show_recovery_warning(
     br_type: str,
     content: str,
     subheader: str | None = None,
     button: str = "TRY AGAIN",
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[ui.UiResult]:
-    await raise_if_not_confirmed(
+    return interact(
-        interact(
+        trezorui2.show_warning(
-            RustLayout(
+            title=content,
-                trezorui2.show_warning(
+            description=subheader or "",
-                    title=content,
+            button=button.upper(),
-                    description=subheader or "",
+            allow_cancel=False,
-                    button=button.upper(),
+        ),
-                    allow_cancel=False,
+        br_type,
-                )
+        br_code,
-            ),
-            br_type,
-            br_code,
-        )
     )

core/src/trezor/ui/layouts/tt/reset.py

@@ -2,14 +2,11 @@ from typing import TYPE_CHECKING
 
 import trezorui2
 from trezor.enums import ButtonRequestType
-from trezor.wire import ActionCancelled
-from trezor.wire.context import wait as ctx_wait
 
 from ..common import interact
-from . import RustLayout, raise_if_not_confirmed
 
 if TYPE_CHECKING:
-    from typing import Callable, Sequence
+    from typing import Awaitable, Callable, Sequence
 
     from trezor.enums import BackupType
 
@@ -55,18 +52,14 @@ async def show_share_words(
 
     pages = _split_share_into_pages(share_words)
 
-    result = await interact(
+    await interact(
-        RustLayout(
+        trezorui2.show_share_words(
-            trezorui2.show_share_words(
+            title=title,
-                title=title,
+            pages=pages,
-                pages=pages,
-            ),
         ),
         "backup_words",
         ButtonRequestType.ResetDevice,
     )
-    if result != CONFIRMED:
-        raise ActionCancelled
 
 
 async def select_word(
@@ -90,14 +83,13 @@ async def select_word(
     while len(words) < 3:
         words.append(words[-1])
 
-    result = await ctx_wait(
+    result = await interact(
-        RustLayout(
+        trezorui2.select_word(
-            trezorui2.select_word(
+            title=title,
-                title=title,
+            description=f"Select word {checked_index + 1} of {count}:",
-                description=f"Select word {checked_index + 1} of {count}:",
+            words=(words[0], words[1], words[2]),
-                words=(words[0], words[1], words[2]),
+        ),
-            )
+        None,
-        )
     )
     if __debug__ and isinstance(result, str):
         return result
@@ -124,20 +116,16 @@ async def slip39_show_checklist(step: int, backup_type: BackupType) -> None:
         )
     )
 
-    result = await interact(
+    await interact(
-        RustLayout(
+        trezorui2.show_checklist(
-            trezorui2.show_checklist(
+            title="BACKUP CHECKLIST",
-                title="BACKUP CHECKLIST",
+            button="CONTINUE",
-                button="CONTINUE",
+            active=step,
-                active=step,
+            items=items,
-                items=items,
-            )
         ),
         "slip39_checklist",
         ButtonRequestType.ResetDevice,
     )
-    if result != CONFIRMED:
-        raise ActionCancelled
 
 
 async def _prompt_number(
@@ -149,14 +137,12 @@ async def _prompt_number(
     max_count: int,
     br_name: str,
 ) -> int:
-    num_input = RustLayout(
+    num_input = trezorui2.request_number(
-        trezorui2.request_number(
+        title=title.upper(),
-            title=title.upper(),
+        description=description,
-            description=description,
+        count=count,
-            count=count,
+        min_count=min_count,
-            min_count=min_count,
+        max_count=max_count,
-            max_count=max_count,
-        )
     )
 
     while True:
@@ -164,26 +150,26 @@ async def _prompt_number(
             num_input,
             br_name,
             ButtonRequestType.ResetDevice,
+            raise_on_cancel=None,
         )
         if __debug__:
             if not isinstance(result, tuple):
                 # DebugLink currently can't send number of shares and it doesn't
                 # change the counter either so just use the initial value.
-                result = (result, count)
+                result = result, count
         status, value = result
 
         if status == CONFIRMED:
             assert isinstance(value, int)
             return value
 
-        await ctx_wait(
+        await interact(
-            RustLayout(
+            trezorui2.show_simple(
-                trezorui2.show_simple(
+                title=None, description=info(value), button="OK, I UNDERSTAND"
-                    title=None, description=info(value), button="OK, I UNDERSTAND"
+            ),
-                )
+            None,
-            )
+            raise_on_cancel=None,
         )
-        num_input.request_complete_repaint()
 
 
 async def slip39_prompt_threshold(
@@ -306,7 +292,7 @@ async def slip39_advanced_prompt_group_threshold(num_of_groups: int) -> int:
     )
 
 
-async def show_warning_backup(slip39: bool) -> None:
+def show_warning_backup(slip39: bool) -> Awaitable[trezorui2.UiResult]:
     if slip39:
         description = (
             "Never make a digital copy of your shares and never upload them online."
@@ -315,46 +301,38 @@ async def show_warning_backup(slip39: bool) -> None:
         description = (
             "Never make a digital copy of your seed and never upload it online."
         )
-    result = await interact(
+    return interact(
-        RustLayout(
+        trezorui2.show_info(
-            trezorui2.show_info(
+            title=description,
-                title=description,
+            button="OK, I UNDERSTAND",
-                button="OK, I UNDERSTAND",
+            allow_cancel=False,
-                allow_cancel=False,
-            )
         ),
         "backup_warning",
         ButtonRequestType.ResetDevice,
     )
-    if result != CONFIRMED:
-        raise ActionCancelled
 
 
-async def show_success_backup() -> None:
+def show_success_backup() -> Awaitable[trezorui2.UiResult]:
     from . import show_success
 
     text = "Use your backup when you need to recover your wallet."
-    await show_success("success_backup", text, "Your backup is done.")
+    return show_success("success_backup", text, "Your backup is done.")
 
 
-async def show_reset_warning(
+def show_reset_warning(
     br_type: str,
     content: str,
     subheader: str | None = None,
     button: str = "TRY AGAIN",
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[trezorui2.UiResult]:
-    await raise_if_not_confirmed(
+    return interact(
-        interact(
+        trezorui2.show_warning(
-            RustLayout(
+            title=subheader or "",
-                trezorui2.show_warning(
+            description=content,
-                    title=subheader or "",
+            button=button.upper(),
-                    description=content,
+            allow_cancel=False,
-                    button=button.upper(),
+        ),
-                    allow_cancel=False,
+        br_type,
-                )
+        br_code,
-            ),
-            br_type,
-            br_code,
-        )
     )

python/src/trezorlib/debuglink.py

@@ -19,6 +19,7 @@ import logging
 import re
 import textwrap
 import time
+from contextlib import contextmanager
 from copy import deepcopy
 from datetime import datetime
 from enum import IntEnum
@@ -38,19 +39,20 @@ from typing import (
     Tuple,
     Type,
     Union,
-    overload,
 )
 
 from mnemonic import Mnemonic
-from typing_extensions import Literal
 
 from . import mapping, messages, protobuf
 from .client import TrezorClient
 from .exceptions import TrezorFailure
 from .log import DUMP_BYTES
+from .messages import DebugWaitType
 from .tools import expect
 
 if TYPE_CHECKING:
+    from typing_extensions import Protocol
+
     from .messages import PinMatrixRequestType
     from .transport import Transport
 
@@ -60,6 +62,15 @@ if TYPE_CHECKING:
 
     AnyDict = Dict[str, Any]
 
+    class InputFunc(Protocol):
+        def __call__(
+            self,
+            hold_ms: Optional[int] = None,
+            wait: Optional[bool] = None,
+        ) -> "LayoutContent":
+            ...
+
+
 EXPECTED_RESPONSES_CONTEXT_LINES = 3
 
 LOG = logging.getLogger(__name__)
@@ -361,6 +372,29 @@ def multipage_content(layouts: List[LayoutContent]) -> str:
     return "".join(layout.text_content() for layout in layouts)
 
 
+def _make_input_func(
+    button: Optional[messages.DebugButton] = None,
+    physical_button: Optional[messages.DebugPhysicalButton] = None,
+    swipe: Optional[messages.DebugSwipeDirection] = None,
+) -> "InputFunc":
+    decision = messages.DebugLinkDecision(
+        button=button,
+        physical_button=physical_button,
+        swipe=swipe,
+    )
+
+    def input_func(
+        self: "DebugLink",
+        hold_ms: Optional[int] = None,
+        wait: Optional[bool] = None,
+    ) -> LayoutContent:
+        __tracebackhide__ = True  # for pytest # pylint: disable=W0612
+        decision.hold_ms = hold_ms
+        return self._decision(decision, wait=wait)
+
+    return input_func  # type: ignore [Parameter name mismatch]
+
+
 class DebugLink:
     def __init__(self, transport: "Transport", auto_interact: bool = True) -> None:
         self.transport = transport
@@ -375,7 +409,6 @@ class DebugLink:
         self.screenshot_recording_dir: Optional[str] = None
 
         # For T1 screenshotting functionality in DebugUI
-        self.t1_take_screenshots = False
         self.t1_screenshot_directory: Optional[Path] = None
         self.t1_screenshot_counter = 0
 
@@ -383,6 +416,11 @@ class DebugLink:
         self.screen_text_file: Optional[Path] = None
         self.last_screen_content = ""
 
+        self.waiting_for_layout_change = False
+        self.layout_dirty = True
+
+        self.input_wait_type = DebugWaitType.IMMEDIATE
+
     @property
     def legacy_ui(self) -> bool:
         """Differences between UI1 and UI2."""
@@ -404,7 +442,12 @@ class DebugLink:
     def close(self) -> None:
         self.transport.end_session()
 
-    def _call(self, msg: protobuf.MessageType, nowait: bool = False) -> Any:
+    def _write(self, msg: protobuf.MessageType) -> None:
+        if self.waiting_for_layout_change:
+            raise RuntimeError(
+                "Debuglink is unavailable while waiting for layout change."
+            )
+
         LOG.debug(
             f"sending message: {msg.__class__.__name__}",
             extra={"protobuf": msg},
@@ -415,13 +458,12 @@ class DebugLink:
             f"encoded as type {msg_type} ({len(msg_bytes)} bytes): {msg_bytes.hex()}",
         )
         self.transport.write(msg_type, msg_bytes)
-        if nowait:
-            return None
 
+    def _read(self) -> protobuf.MessageType:
         ret_type, ret_bytes = self.transport.read()
         LOG.log(
             DUMP_BYTES,
-            f"received type {msg_type} ({len(msg_bytes)} bytes): {msg_bytes.hex()}",
+            f"received type {ret_type} ({len(ret_bytes)} bytes): {ret_bytes.hex()}",
         )
         msg = self.mapping.decode(ret_type, ret_bytes)
         LOG.debug(
@@ -430,11 +472,20 @@ class DebugLink:
         )
         return msg
 
-    def state(self) -> messages.DebugLinkState:
+    def _call(self, msg: protobuf.MessageType) -> Any:
-        return self._call(messages.DebugLinkGetState())
+        self._write(msg)
+        return self._read()
+
+    def state(
+        self, wait_type: DebugWaitType = DebugWaitType.CURRENT_LAYOUT
+    ) -> messages.DebugLinkState:
+        result = self._call(messages.DebugLinkGetState(wait_layout=wait_type))
+        if isinstance(result, messages.Failure):
+            raise TrezorFailure(result)
+        return result
 
     def read_layout(self) -> LayoutContent:
-        return LayoutContent(self.state().tokens or [])
+        return LayoutContent(self.state().tokens)
 
     def wait_layout(self, wait_for_external_change: bool = False) -> LayoutContent:
         # Next layout change will be caused by external event
@@ -445,11 +496,38 @@ class DebugLink:
         if wait_for_external_change:
             self.reset_debug_events()
 
-        obj = self._call(messages.DebugLinkGetState(wait_layout=True))
+        obj = self._call(
+            messages.DebugLinkGetState(wait_layout=DebugWaitType.NEXT_LAYOUT)
+        )
+        self.layout_dirty = True
         if isinstance(obj, messages.Failure):
             raise TrezorFailure(obj)
         return LayoutContent(obj.tokens)
 
+    @contextmanager
+    def wait_for_layout_change(self) -> Iterator[LayoutContent]:
+        # set up a dummy layout content object to be yielded
+        layout_content = LayoutContent(
+            ["DUMMY CONTENT, WAIT UNTIL THE END OF THE BLOCK :("]
+        )
+
+        # send GetState without waiting for reply
+        self._write(messages.DebugLinkGetState(wait_layout=DebugWaitType.NEXT_LAYOUT))
+
+        # allow the block to proceed
+        self.waiting_for_layout_change = True
+        try:
+            yield layout_content
+        finally:
+            self.waiting_for_layout_change = False
+
+        # wait for the reply
+        resp = self._read()
+        assert isinstance(resp, messages.DebugLinkState)
+
+        # replace contents of the yielded object with the new thing
+        layout_content.__init__(resp.tokens)
+
     def reset_debug_events(self) -> None:
         # Only supported on TT and above certain version
         if self.model in ("T", "Safe 3") and not self.legacy_debug:
@@ -493,56 +571,102 @@ class DebugLink:
         state = self._call(messages.DebugLinkGetState(wait_word_list=True))
         return state.reset_word
 
-    def input(
+    def _decision(
-        self,
+        self, decision: messages.DebugLinkDecision, wait: Optional[bool] = None
-        word: Optional[str] = None,
+    ) -> LayoutContent:
-        button: Optional[messages.DebugButton] = None,
+        """Send a debuglink decision and returns the resulting layout.
-        physical_button: Optional[messages.DebugPhysicalButton] = None,
+
-        swipe: Optional[messages.DebugSwipeDirection] = None,
+        If hold_ms is set, an additional 200ms is added to account for processing
-        x: Optional[int] = None,
+        delays. (This is needed for hold-to-confirm to trigger reliably.)
-        y: Optional[int] = None,
+
-        wait: Optional[bool] = None,
+        If `wait` is unset, the current wait mode is used:
-        hold_ms: Optional[int] = None,
+
-    ) -> Optional[LayoutContent]:
+        - when in normal tests, IMMEDIATE, which never deadlocks the device, but may
+          return an empty layout in case the next one didn't come up immediately. (E.g.,
+          in SignTx flow, the device is waiting for more TxRequest/TxAck exchanges
+          before showing the next UI layout.)
+        - when in tests running through a `DeviceHandler`, CURRENT_LAYOUT, which waits
+          for the next layout to come up. The assumption is that wirelink is
+          communicating on another thread and won't be blocked by waiting on debuglink.
+
+        Force waiting for the layout by setting `wait=True`. Force not waiting by
+        setting `wait=False` -- useful when, e.g., you are causing the next layout to be
+        deliberately delayed.
+        """
         if not self.allow_interactions:
-            return None
+            return self.wait_layout()
 
-        args = sum(a is not None for a in (word, button, physical_button, swipe, x))
+        if decision.hold_ms is not None:
-        if args != 1:
+            decision.hold_ms += 200
-            raise ValueError(
-                "Invalid input - must use one of word, button, physical_button, swipe, click(x,y)"
-            )
 
-        decision = messages.DebugLinkDecision(
+        self._write(decision)
-            button=button,
+        self.layout_dirty = True
-            physical_button=physical_button,
+        if wait is True:
-            swipe=swipe,
+            wait_type = DebugWaitType.CURRENT_LAYOUT
-            input=word,
+        elif wait is False:
-            x=x,
+            wait_type = DebugWaitType.IMMEDIATE
-            y=y,
+        else:
-            wait=wait,
+            wait_type = self.input_wait_type
-            hold_ms=hold_ms,
+        return self.snapshot(wait_type)
+
+    press_yes = _make_input_func(button=messages.DebugButton.YES)
+    """Confirm current layout. See `_decision` for more details."""
+    press_no = _make_input_func(button=messages.DebugButton.NO)
+    """Reject current layout. See `_decision` for more details."""
+    press_info = _make_input_func(button=messages.DebugButton.INFO)
+    """Trigger the Info action. See `_decision` for more details."""
+    swipe_up = _make_input_func(swipe=messages.DebugSwipeDirection.UP)
+    """Swipe up. See `_decision` for more details."""
+    swipe_down = _make_input_func(swipe=messages.DebugSwipeDirection.DOWN)
+    """Swipe down. See `_decision` for more details."""
+    swipe_right = _make_input_func(swipe=messages.DebugSwipeDirection.RIGHT)
+    """Swipe right. See `_decision` for more details."""
+    swipe_left = _make_input_func(swipe=messages.DebugSwipeDirection.LEFT)
+    """Swipe left. See `_decision` for more details."""
+    press_left = _make_input_func(physical_button=messages.DebugPhysicalButton.LEFT_BTN)
+    """Press left button. See `_decision` for more details."""
+    press_middle = _make_input_func(
+        physical_button=messages.DebugPhysicalButton.MIDDLE_BTN
+    )
+    """Press middle button. See `_decision` for more details."""
+    press_right = _make_input_func(
+        physical_button=messages.DebugPhysicalButton.RIGHT_BTN
+    )
+    """Press right button. See `_decision` for more details."""
+
+    def input(self, word: str, wait: Optional[bool] = None) -> LayoutContent:
+        """Send text input to the device. See `_decision` for more details."""
+        return self._decision(messages.DebugLinkDecision(input=word), wait)
+
+    def click(
+        self,
+        click: Tuple[int, int],
+        hold_ms: Optional[int] = None,
+        wait: Optional[bool] = None,
+    ) -> LayoutContent:
+        """Send a click to the device. See `_decision` for more details."""
+        x, y = click
+        return self._decision(
+            messages.DebugLinkDecision(x=x, y=y, hold_ms=hold_ms), wait
         )
 
-        ret = self._call(decision, nowait=not wait)
+    def snapshot(
-        if ret is not None:
+        self, wait_type: DebugWaitType = DebugWaitType.IMMEDIATE
-            return LayoutContent(ret.tokens)
+    ) -> LayoutContent:
+        """Save text and image content of the screen to relevant directories."""
+        # take the snapshot
+        state = self.state(wait_type)
+        layout = LayoutContent(state.tokens)
 
-        # Getting the current screen after the (nowait) decision
+        if state.tokens and self.layout_dirty:
-        self.save_current_screen_if_relevant(wait=False)
+            # save it, unless we already did or unless it's empty
+            self.save_debug_screen(layout.visible_screen())
+            if state.layout is not None:
+                self.save_screenshot(state.layout)
+            self.layout_dirty = False
 
-        return None
+        # return the layout
-
+        return layout
-    def save_current_screen_if_relevant(self, wait: bool = True) -> None:
-        """Optionally saving the textual screen output."""
-        if self.screen_text_file is None:
-            return
-
-        if wait:
-            layout = self.wait_layout()
-        else:
-            layout = self.read_layout()
-        self.save_debug_screen(layout.visible_screen())
 
     def save_debug_screen(self, screen_content: str) -> None:
         if self.screen_text_file is None:
@@ -561,139 +685,8 @@ class DebugLink:
             f.write(screen_content)
             f.write("\n" + 80 * "/" + "\n")
 
-    # Type overloads below make sure that when we supply `wait=True` into functions,
-    # they will always return `LayoutContent` and we do not need to assert `is not None`.
-
-    @overload
-    def click(self, click: Tuple[int, int]) -> None:
-        ...
-
-    @overload
-    def click(self, click: Tuple[int, int], wait: Literal[True]) -> LayoutContent:
-        ...
-
-    def click(
-        self, click: Tuple[int, int], wait: bool = False
-    ) -> Optional[LayoutContent]:
-        x, y = click
-        return self.input(x=x, y=y, wait=wait)
-
-    # Made into separate function as `hold_ms: Optional[int]` in `click`
-    # was causing problems with @overload
-    def click_hold(
-        self, click: Tuple[int, int], hold_ms: int
-    ) -> Optional[LayoutContent]:
-        x, y = click
-        return self.input(x=x, y=y, hold_ms=hold_ms, wait=True)
-
-    def press_yes(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(button=messages.DebugButton.YES, wait=wait)
-
-    def press_no(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(button=messages.DebugButton.NO, wait=wait)
-
-    def press_info(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(button=messages.DebugButton.INFO, wait=wait)
-
-    def swipe_up(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(swipe=messages.DebugSwipeDirection.UP, wait=wait)
-
-    def swipe_down(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(swipe=messages.DebugSwipeDirection.DOWN, wait=wait)
-
-    @overload
-    def swipe_right(self) -> None:
-        ...
-
-    @overload
-    def swipe_right(self, wait: Literal[True]) -> LayoutContent:
-        ...
-
-    def swipe_right(self, wait: bool = False) -> Union[LayoutContent, None]:
-        return self.input(swipe=messages.DebugSwipeDirection.RIGHT, wait=wait)
-
-    @overload
-    def swipe_left(self) -> None:
-        ...
-
-    @overload
-    def swipe_left(self, wait: Literal[True]) -> LayoutContent:
-        ...
-
-    def swipe_left(self, wait: bool = False) -> Union[LayoutContent, None]:
-        return self.input(swipe=messages.DebugSwipeDirection.LEFT, wait=wait)
-
-    @overload
-    def press_left(self) -> None:
-        ...
-
-    @overload
-    def press_left(self, wait: Literal[True]) -> LayoutContent:
-        ...
-
-    def press_left(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(
-            physical_button=messages.DebugPhysicalButton.LEFT_BTN, wait=wait
-        )
-
-    @overload
-    def press_middle(self) -> None:
-        ...
-
-    @overload
-    def press_middle(self, wait: Literal[True]) -> LayoutContent:
-        ...
-
-    def press_middle(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(
-            physical_button=messages.DebugPhysicalButton.MIDDLE_BTN, wait=wait
-        )
-
-    def press_middle_htc(
-        self, hold_ms: int, extra_ms: int = 200
-    ) -> Optional[LayoutContent]:
-        return self.press_htc(
-            button=messages.DebugPhysicalButton.MIDDLE_BTN,
-            hold_ms=hold_ms,
-            extra_ms=extra_ms,
-        )
-
-    @overload
-    def press_right(self) -> None:
-        ...
-
-    @overload
-    def press_right(self, wait: Literal[True]) -> LayoutContent:
-        ...
-
-    def press_right(self, wait: bool = False) -> Optional[LayoutContent]:
-        return self.input(
-            physical_button=messages.DebugPhysicalButton.RIGHT_BTN, wait=wait
-        )
-
-    def press_right_htc(
-        self, hold_ms: int, extra_ms: int = 200
-    ) -> Optional[LayoutContent]:
-        return self.press_htc(
-            button=messages.DebugPhysicalButton.RIGHT_BTN,
-            hold_ms=hold_ms,
-            extra_ms=extra_ms,
-        )
-
-    def press_htc(
-        self, button: messages.DebugPhysicalButton, hold_ms: int, extra_ms: int = 200
-    ) -> Optional[LayoutContent]:
-        hold_ms = hold_ms + extra_ms  # safety margin
-        result = self.input(
-            physical_button=button,
-            hold_ms=hold_ms,
-        )
-        # sleeping little longer for UI to update
-        time.sleep(hold_ms / 1000 + 0.1)
-        return result
-
     def stop(self) -> None:
-        self._call(messages.DebugLinkStop(), nowait=True)
+        self._write(messages.DebugLinkStop())
 
     def reseed(self, value: int) -> protobuf.MessageType:
         return self._call(messages.DebugLinkReseedRandom(value=value))
@@ -727,44 +720,35 @@ class DebugLink:
         return self._call(messages.DebugLinkMemoryRead(address=address, length=length))
 
     def memory_write(self, address: int, memory: bytes, flash: bool = False) -> None:
-        self._call(
+        self._write(
-            messages.DebugLinkMemoryWrite(address=address, memory=memory, flash=flash),
+            messages.DebugLinkMemoryWrite(address=address, memory=memory, flash=flash)
-            nowait=True,
         )
 
     def flash_erase(self, sector: int) -> None:
-        self._call(messages.DebugLinkFlashErase(sector=sector), nowait=True)
+        self._write(messages.DebugLinkFlashErase(sector=sector))
 
     @expect(messages.Success)
     def erase_sd_card(self, format: bool = True) -> messages.Success:
         return self._call(messages.DebugLinkEraseSdCard(format=format))
 
-    def take_t1_screenshot_if_relevant(self) -> None:
+    def save_screenshot(self, data: bytes) -> None:
-        """Conditionally take screenshots on T1.
+        if self.t1_screenshot_directory is None:
+            return
 
-        TT handles them differently, see debuglink.start_recording.
-        """
-        if self.model == "1" and self.t1_take_screenshots:
-            self.save_screenshot_for_t1()
-
-    def save_screenshot_for_t1(self) -> None:
         from PIL import Image
 
-        layout = self.state().layout
+        assert len(data) == 128 * 64 // 8
-        assert layout is not None
-        assert len(layout) == 128 * 64 // 8
 
         pixels: List[int] = []
         for byteline in range(64 // 8):
             offset = byteline * 128
-            row = layout[offset : offset + 128]
+            row = data[offset : offset + 128]
             for bit in range(8):
                 pixels.extend(bool(px & (1 << bit)) for px in row)
 
         im = Image.new("1", (128, 64))
         im.putdata(pixels[::-1])
 
-        assert self.t1_screenshot_directory is not None
         img_location = (
             self.t1_screenshot_directory / f"{self.t1_screenshot_counter:04d}.png"
         )
@@ -772,6 +756,9 @@ class DebugLink:
         self.t1_screenshot_counter += 1
 
 
+del _make_input_func
+
+
 class NullDebugLink(DebugLink):
     def __init__(self) -> None:
         # Ignoring type error as self.transport will not be touched while using NullDebugLink
@@ -810,15 +797,9 @@ class DebugUI:
         ] = None
 
     def button_request(self, br: messages.ButtonRequest) -> None:
-        self.debuglink.take_t1_screenshot_if_relevant()
+        self.debuglink.snapshot()
 
         if self.input_flow is None:
-            # Only calling screen-saver when not in input-flow
-            # as it collides with wait-layout of input flows.
-            # All input flows call debuglink.input(), so
-            # recording their screens that way (as well as
-            # possible swipes below).
-            self.debuglink.save_current_screen_if_relevant(wait=True)
             if br.code == messages.ButtonRequestType.PinEntry:
                 self.debuglink.input(self.get_pin())
             else:
@@ -837,7 +818,7 @@ class DebugUI:
                 self.input_flow = self.INPUT_FLOW_DONE
 
     def get_pin(self, code: Optional["PinMatrixRequestType"] = None) -> str:
-        self.debuglink.take_t1_screenshot_if_relevant()
+        self.debuglink.snapshot()
 
         if self.pins is None:
             raise RuntimeError("PIN requested but no sequence was configured")
@@ -848,7 +829,7 @@ class DebugUI:
             raise AssertionError("PIN sequence ended prematurely")
 
     def get_passphrase(self, available_on_device: bool) -> str:
-        self.debuglink.take_t1_screenshot_if_relevant()
+        self.debuglink.snapshot()
         return self.passphrase
 
 

python/src/trezorlib/messages.py

@@ -509,6 +509,12 @@ class DebugPhysicalButton(IntEnum):
     RIGHT_BTN = 2
 
 
+class DebugWaitType(IntEnum):
+    IMMEDIATE = 0
+    NEXT_LAYOUT = 1
+    CURRENT_LAYOUT = 2
+
+
 class EthereumDefinitionType(IntEnum):
     NETWORK = 0
     TOKEN = 1
@@ -3908,7 +3914,7 @@ class DebugLinkGetState(protobuf.MessageType):
     FIELDS = {
         1: protobuf.Field("wait_word_list", "bool", repeated=False, required=False, default=None),
         2: protobuf.Field("wait_word_pos", "bool", repeated=False, required=False, default=None),
-        3: protobuf.Field("wait_layout", "bool", repeated=False, required=False, default=None),
+        3: protobuf.Field("wait_layout", "DebugWaitType", repeated=False, required=False, default=DebugWaitType.IMMEDIATE),
     }
 
     def __init__(
@@ -3916,7 +3922,7 @@ class DebugLinkGetState(protobuf.MessageType):
         *,
         wait_word_list: Optional["bool"] = None,
         wait_word_pos: Optional["bool"] = None,
-        wait_layout: Optional["bool"] = None,
+        wait_layout: Optional["DebugWaitType"] = DebugWaitType.IMMEDIATE,
     ) -> None:
         self.wait_word_list = wait_word_list
         self.wait_word_pos = wait_word_pos

python/src/trezorlib/protobuf.py

@@ -509,12 +509,17 @@ def format_message(
         return printable / len(bytes) > 0.8
 
     def pformat(name: str, value: Any, indent: int) -> str:
+        from . import messages
+
         level = sep * indent
         leadin = sep * (indent + 1)
 
         if isinstance(value, MessageType):
             return format_message(value, indent, sep)
 
+        if isinstance(pb, messages.DebugLinkState) and name == "tokens":
+            return "".join(value)
+
         if isinstance(value, list):
             # short list of simple values
             if not value or all(isinstance(x, int) for x in value):

rust/trezor-client/src/protos/generated/messages_debug.rs

@@ -1127,7 +1127,7 @@ pub struct DebugLinkGetState {
     // @@protoc_insertion_point(field:hw.trezor.messages.debug.DebugLinkGetState.wait_word_pos)
     pub wait_word_pos: ::std::option::Option<bool>,
     // @@protoc_insertion_point(field:hw.trezor.messages.debug.DebugLinkGetState.wait_layout)
-    pub wait_layout: ::std::option::Option<bool>,
+    pub wait_layout: ::std::option::Option<::protobuf::EnumOrUnknown<debug_link_get_state::DebugWaitType>>,
     // special fields
     // @@protoc_insertion_point(special_field:hw.trezor.messages.debug.DebugLinkGetState.special_fields)
     pub special_fields: ::protobuf::SpecialFields,
@@ -1182,10 +1182,13 @@ impl DebugLinkGetState {
         self.wait_word_pos = ::std::option::Option::Some(v);
     }
 
-    // optional bool wait_layout = 3;
+    // optional .hw.trezor.messages.debug.DebugLinkGetState.DebugWaitType wait_layout = 3;
 
-    pub fn wait_layout(&self) -> bool {
+    pub fn wait_layout(&self) -> debug_link_get_state::DebugWaitType {
-        self.wait_layout.unwrap_or(false)
+        match self.wait_layout {
+            Some(e) => e.enum_value_or(debug_link_get_state::DebugWaitType::IMMEDIATE),
+            None => debug_link_get_state::DebugWaitType::IMMEDIATE,
+        }
     }
 
     pub fn clear_wait_layout(&mut self) {
@@ -1197,8 +1200,8 @@ impl DebugLinkGetState {
     }
 
     // Param is passed by value, moved
-    pub fn set_wait_layout(&mut self, v: bool) {
+    pub fn set_wait_layout(&mut self, v: debug_link_get_state::DebugWaitType) {
-        self.wait_layout = ::std::option::Option::Some(v);
+        self.wait_layout = ::std::option::Option::Some(::protobuf::EnumOrUnknown::new(v));
     }
 
     fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData {
@@ -1244,7 +1247,7 @@ impl ::protobuf::Message for DebugLinkGetState {
                     self.wait_word_pos = ::std::option::Option::Some(is.read_bool()?);
                 },
                 24 => {
-                    self.wait_layout = ::std::option::Option::Some(is.read_bool()?);
+                    self.wait_layout = ::std::option::Option::Some(is.read_enum_or_unknown()?);
                 },
                 tag => {
                     ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?;
@@ -1265,7 +1268,7 @@ impl ::protobuf::Message for DebugLinkGetState {
             my_size += 1 + 1;
         }
         if let Some(v) = self.wait_layout {
-            my_size += 1 + 1;
+            my_size += ::protobuf::rt::int32_size(3, v.value());
         }
         my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields());
         self.special_fields.cached_size().set(my_size as u32);
@@ -1280,7 +1283,7 @@ impl ::protobuf::Message for DebugLinkGetState {
             os.write_bool(2, v)?;
         }
         if let Some(v) = self.wait_layout {
-            os.write_bool(3, v)?;
+            os.write_enum(3, ::protobuf::EnumOrUnknown::value(&v))?;
         }
         os.write_unknown_fields(self.special_fields.unknown_fields())?;
         ::std::result::Result::Ok(())
@@ -1333,6 +1336,76 @@ impl ::protobuf::reflect::ProtobufValue for DebugLinkGetState {
     type RuntimeType = ::protobuf::reflect::rt::RuntimeTypeMessage<Self>;
 }
 
+/// Nested message and enums of message `DebugLinkGetState`
+pub mod debug_link_get_state {
+    #[derive(Clone,Copy,PartialEq,Eq,Debug,Hash)]
+    // @@protoc_insertion_point(enum:hw.trezor.messages.debug.DebugLinkGetState.DebugWaitType)
+    pub enum DebugWaitType {
+        // @@protoc_insertion_point(enum_value:hw.trezor.messages.debug.DebugLinkGetState.DebugWaitType.IMMEDIATE)
+        IMMEDIATE = 0,
+        // @@protoc_insertion_point(enum_value:hw.trezor.messages.debug.DebugLinkGetState.DebugWaitType.NEXT_LAYOUT)
+        NEXT_LAYOUT = 1,
+        // @@protoc_insertion_point(enum_value:hw.trezor.messages.debug.DebugLinkGetState.DebugWaitType.CURRENT_LAYOUT)
+        CURRENT_LAYOUT = 2,
+    }
+
+    impl ::protobuf::Enum for DebugWaitType {
+        const NAME: &'static str = "DebugWaitType";
+
+        fn value(&self) -> i32 {
+            *self as i32
+        }
+
+        fn from_i32(value: i32) -> ::std::option::Option<DebugWaitType> {
+            match value {
+                0 => ::std::option::Option::Some(DebugWaitType::IMMEDIATE),
+                1 => ::std::option::Option::Some(DebugWaitType::NEXT_LAYOUT),
+                2 => ::std::option::Option::Some(DebugWaitType::CURRENT_LAYOUT),
+                _ => ::std::option::Option::None
+            }
+        }
+
+        fn from_str(str: &str) -> ::std::option::Option<DebugWaitType> {
+            match str {
+                "IMMEDIATE" => ::std::option::Option::Some(DebugWaitType::IMMEDIATE),
+                "NEXT_LAYOUT" => ::std::option::Option::Some(DebugWaitType::NEXT_LAYOUT),
+                "CURRENT_LAYOUT" => ::std::option::Option::Some(DebugWaitType::CURRENT_LAYOUT),
+                _ => ::std::option::Option::None
+            }
+        }
+
+        const VALUES: &'static [DebugWaitType] = &[
+            DebugWaitType::IMMEDIATE,
+            DebugWaitType::NEXT_LAYOUT,
+            DebugWaitType::CURRENT_LAYOUT,
+        ];
+    }
+
+    impl ::protobuf::EnumFull for DebugWaitType {
+        fn enum_descriptor() -> ::protobuf::reflect::EnumDescriptor {
+            static descriptor: ::protobuf::rt::Lazy<::protobuf::reflect::EnumDescriptor> = ::protobuf::rt::Lazy::new();
+            descriptor.get(|| super::file_descriptor().enum_by_package_relative_name("DebugLinkGetState.DebugWaitType").unwrap()).clone()
+        }
+
+        fn descriptor(&self) -> ::protobuf::reflect::EnumValueDescriptor {
+            let index = *self as usize;
+            Self::enum_descriptor().value_by_index(index)
+        }
+    }
+
+    impl ::std::default::Default for DebugWaitType {
+        fn default() -> Self {
+            DebugWaitType::IMMEDIATE
+        }
+    }
+
+    impl DebugWaitType {
+        pub(in super) fn generated_enum_descriptor_data() -> ::protobuf::reflect::GeneratedEnumDescriptorData {
+            ::protobuf::reflect::GeneratedEnumDescriptorData::new::<DebugWaitType>("DebugLinkGetState.DebugWaitType")
+        }
+    }
+}
+
 // @@protoc_insertion_point(message:hw.trezor.messages.debug.DebugLinkState)
 #[derive(PartialEq,Clone,Default,Debug)]
 pub struct DebugLinkState {
@@ -3457,52 +3530,56 @@ impl ::protobuf::reflect::ProtobufValue for DebugLinkResetDebugEvents {
 static file_descriptor_proto_data: &'static [u8] = b"\
     \n\x14messages-debug.proto\x12\x18hw.trezor.messages.debug\x1a\x0emessag\
     es.proto\x1a\x15messages-common.proto\x1a\x19messages-management.proto\"\
-    \xb0\x04\n\x11DebugLinkDecision\x12O\n\x06button\x18\x01\x20\x01(\x0e27.\
+    \xb4\x04\n\x11DebugLinkDecision\x12O\n\x06button\x18\x01\x20\x01(\x0e27.\
     hw.trezor.messages.debug.DebugLinkDecision.DebugButtonR\x06button\x12U\n\
     \x05swipe\x18\x02\x20\x01(\x0e2?.hw.trezor.messages.debug.DebugLinkDecis\
     ion.DebugSwipeDirectionR\x05swipe\x12\x14\n\x05input\x18\x03\x20\x01(\tR\
     \x05input\x12\x0c\n\x01x\x18\x04\x20\x01(\rR\x01x\x12\x0c\n\x01y\x18\x05\
-    \x20\x01(\rR\x01y\x12\x12\n\x04wait\x18\x06\x20\x01(\x08R\x04wait\x12\
+    \x20\x01(\rR\x01y\x12\x16\n\x04wait\x18\x06\x20\x01(\x08R\x04waitB\x02\
-    \x17\n\x07hold_ms\x18\x07\x20\x01(\rR\x06holdMs\x12h\n\x0fphysical_butto\
+    \x18\x01\x12\x17\n\x07hold_ms\x18\x07\x20\x01(\rR\x06holdMs\x12h\n\x0fph\
-    n\x18\x08\x20\x01(\x0e2?.hw.trezor.messages.debug.DebugLinkDecision.Debu\
+    ysical_button\x18\x08\x20\x01(\x0e2?.hw.trezor.messages.debug.DebugLinkD\
-    gPhysicalButtonR\x0ephysicalButton\"<\n\x13DebugSwipeDirection\x12\x06\n\
+    ecision.DebugPhysicalButtonR\x0ephysicalButton\"<\n\x13DebugSwipeDirecti\
-    \x02UP\x10\0\x12\x08\n\x04DOWN\x10\x01\x12\x08\n\x04LEFT\x10\x02\x12\t\n\
+    on\x12\x06\n\x02UP\x10\0\x12\x08\n\x04DOWN\x10\x01\x12\x08\n\x04LEFT\x10\
-    \x05RIGHT\x10\x03\"(\n\x0bDebugButton\x12\x06\n\x02NO\x10\0\x12\x07\n\
+    \x02\x12\t\n\x05RIGHT\x10\x03\"(\n\x0bDebugButton\x12\x06\n\x02NO\x10\0\
-    \x03YES\x10\x01\x12\x08\n\x04INFO\x10\x02\"B\n\x13DebugPhysicalButton\
+    \x12\x07\n\x03YES\x10\x01\x12\x08\n\x04INFO\x10\x02\"B\n\x13DebugPhysica\
-    \x12\x0c\n\x08LEFT_BTN\x10\0\x12\x0e\n\nMIDDLE_BTN\x10\x01\x12\r\n\tRIGH\
+    lButton\x12\x0c\n\x08LEFT_BTN\x10\0\x12\x0e\n\nMIDDLE_BTN\x10\x01\x12\r\
-    T_BTN\x10\x02\")\n\x0fDebugLinkLayout\x12\x16\n\x06tokens\x18\x01\x20\
+    \n\tRIGHT_BTN\x10\x02\"-\n\x0fDebugLinkLayout\x12\x16\n\x06tokens\x18\
-    \x03(\tR\x06tokens\"-\n\x15DebugLinkReseedRandom\x12\x14\n\x05value\x18\
+    \x01\x20\x03(\tR\x06tokens:\x02\x18\x01\"-\n\x15DebugLinkReseedRandom\
-    \x01\x20\x01(\rR\x05value\"j\n\x15DebugLinkRecordScreen\x12)\n\x10target\
+    \x12\x14\n\x05value\x18\x01\x20\x01(\rR\x05value\"j\n\x15DebugLinkRecord\
-    _directory\x18\x01\x20\x01(\tR\x0ftargetDirectory\x12&\n\rrefresh_index\
+    Screen\x12)\n\x10target_directory\x18\x01\x20\x01(\tR\x0ftargetDirectory\
-    \x18\x02\x20\x01(\r:\x010R\x0crefreshIndex\"~\n\x11DebugLinkGetState\x12\
+    \x12&\n\rrefresh_index\x18\x02\x20\x01(\r:\x010R\x0crefreshIndex\"\x91\
-    $\n\x0ewait_word_list\x18\x01\x20\x01(\x08R\x0cwaitWordList\x12\"\n\rwai\
+    \x02\n\x11DebugLinkGetState\x12(\n\x0ewait_word_list\x18\x01\x20\x01(\
-    t_word_pos\x18\x02\x20\x01(\x08R\x0bwaitWordPos\x12\x1f\n\x0bwait_layout\
+    \x08R\x0cwaitWordListB\x02\x18\x01\x12&\n\rwait_word_pos\x18\x02\x20\x01\
-    \x18\x03\x20\x01(\x08R\nwaitLayout\"\x97\x04\n\x0eDebugLinkState\x12\x16\
+    (\x08R\x0bwaitWordPosB\x02\x18\x01\x12e\n\x0bwait_layout\x18\x03\x20\x01\
-    \n\x06layout\x18\x01\x20\x01(\x0cR\x06layout\x12\x10\n\x03pin\x18\x02\
+    (\x0e29.hw.trezor.messages.debug.DebugLinkGetState.DebugWaitType:\tIMMED\
-    \x20\x01(\tR\x03pin\x12\x16\n\x06matrix\x18\x03\x20\x01(\tR\x06matrix\
+    IATER\nwaitLayout\"C\n\rDebugWaitType\x12\r\n\tIMMEDIATE\x10\0\x12\x0f\n\
-    \x12'\n\x0fmnemonic_secret\x18\x04\x20\x01(\x0cR\x0emnemonicSecret\x129\
+    \x0bNEXT_LAYOUT\x10\x01\x12\x12\n\x0eCURRENT_LAYOUT\x10\x02\"\x97\x04\n\
-    \n\x04node\x18\x05\x20\x01(\x0b2%.hw.trezor.messages.common.HDNodeTypeR\
+    \x0eDebugLinkState\x12\x16\n\x06layout\x18\x01\x20\x01(\x0cR\x06layout\
-    \x04node\x123\n\x15passphrase_protection\x18\x06\x20\x01(\x08R\x14passph\
+    \x12\x10\n\x03pin\x18\x02\x20\x01(\tR\x03pin\x12\x16\n\x06matrix\x18\x03\
-    raseProtection\x12\x1d\n\nreset_word\x18\x07\x20\x01(\tR\tresetWord\x12#\
+    \x20\x01(\tR\x06matrix\x12'\n\x0fmnemonic_secret\x18\x04\x20\x01(\x0cR\
-    \n\rreset_entropy\x18\x08\x20\x01(\x0cR\x0cresetEntropy\x12,\n\x12recove\
+    \x0emnemonicSecret\x129\n\x04node\x18\x05\x20\x01(\x0b2%.hw.trezor.messa\
-    ry_fake_word\x18\t\x20\x01(\tR\x10recoveryFakeWord\x12*\n\x11recovery_wo\
+    ges.common.HDNodeTypeR\x04node\x123\n\x15passphrase_protection\x18\x06\
-    rd_pos\x18\n\x20\x01(\rR\x0frecoveryWordPos\x12$\n\x0ereset_word_pos\x18\
+    \x20\x01(\x08R\x14passphraseProtection\x12\x1d\n\nreset_word\x18\x07\x20\
-    \x0b\x20\x01(\rR\x0cresetWordPos\x12N\n\rmnemonic_type\x18\x0c\x20\x01(\
+    \x01(\tR\tresetWord\x12#\n\rreset_entropy\x18\x08\x20\x01(\x0cR\x0creset\
-    \x0e2).hw.trezor.messages.management.BackupTypeR\x0cmnemonicType\x12\x16\
+    Entropy\x12,\n\x12recovery_fake_word\x18\t\x20\x01(\tR\x10recoveryFakeWo\
-    \n\x06tokens\x18\r\x20\x03(\tR\x06tokens\"\x0f\n\rDebugLinkStop\"P\n\x0c\
+    rd\x12*\n\x11recovery_word_pos\x18\n\x20\x01(\rR\x0frecoveryWordPos\x12$\
-    DebugLinkLog\x12\x14\n\x05level\x18\x01\x20\x01(\rR\x05level\x12\x16\n\
+    \n\x0ereset_word_pos\x18\x0b\x20\x01(\rR\x0cresetWordPos\x12N\n\rmnemoni\
-    \x06bucket\x18\x02\x20\x01(\tR\x06bucket\x12\x12\n\x04text\x18\x03\x20\
+    c_type\x18\x0c\x20\x01(\x0e2).hw.trezor.messages.management.BackupTypeR\
-    \x01(\tR\x04text\"G\n\x13DebugLinkMemoryRead\x12\x18\n\x07address\x18\
+    \x0cmnemonicType\x12\x16\n\x06tokens\x18\r\x20\x03(\tR\x06tokens\"\x0f\n\
-    \x01\x20\x01(\rR\x07address\x12\x16\n\x06length\x18\x02\x20\x01(\rR\x06l\
+    \rDebugLinkStop\"P\n\x0cDebugLinkLog\x12\x14\n\x05level\x18\x01\x20\x01(\
-    ength\")\n\x0fDebugLinkMemory\x12\x16\n\x06memory\x18\x01\x20\x01(\x0cR\
+    \rR\x05level\x12\x16\n\x06bucket\x18\x02\x20\x01(\tR\x06bucket\x12\x12\n\
-    \x06memory\"^\n\x14DebugLinkMemoryWrite\x12\x18\n\x07address\x18\x01\x20\
+    \x04text\x18\x03\x20\x01(\tR\x04text\"G\n\x13DebugLinkMemoryRead\x12\x18\
-    \x01(\rR\x07address\x12\x16\n\x06memory\x18\x02\x20\x01(\x0cR\x06memory\
+    \n\x07address\x18\x01\x20\x01(\rR\x07address\x12\x16\n\x06length\x18\x02\
-    \x12\x14\n\x05flash\x18\x03\x20\x01(\x08R\x05flash\"-\n\x13DebugLinkFlas\
+    \x20\x01(\rR\x06length\")\n\x0fDebugLinkMemory\x12\x16\n\x06memory\x18\
-    hErase\x12\x16\n\x06sector\x18\x01\x20\x01(\rR\x06sector\".\n\x14DebugLi\
+    \x01\x20\x01(\x0cR\x06memory\"^\n\x14DebugLinkMemoryWrite\x12\x18\n\x07a\
-    nkEraseSdCard\x12\x16\n\x06format\x18\x01\x20\x01(\x08R\x06format\",\n\
+    ddress\x18\x01\x20\x01(\rR\x07address\x12\x16\n\x06memory\x18\x02\x20\
-    \x14DebugLinkWatchLayout\x12\x14\n\x05watch\x18\x01\x20\x01(\x08R\x05wat\
+    \x01(\x0cR\x06memory\x12\x14\n\x05flash\x18\x03\x20\x01(\x08R\x05flash\"\
-    ch\"\x1b\n\x19DebugLinkResetDebugEventsB=\n#com.satoshilabs.trezor.lib.p\
+    -\n\x13DebugLinkFlashErase\x12\x16\n\x06sector\x18\x01\x20\x01(\rR\x06se\
-    rotobufB\x12TrezorMessageDebug\x80\xa6\x1d\x01\
+    ctor\".\n\x14DebugLinkEraseSdCard\x12\x16\n\x06format\x18\x01\x20\x01(\
+    \x08R\x06format\"0\n\x14DebugLinkWatchLayout\x12\x14\n\x05watch\x18\x01\
+    \x20\x01(\x08R\x05watch:\x02\x18\x01\"\x1f\n\x19DebugLinkResetDebugEvent\
+    s:\x02\x18\x01B=\n#com.satoshilabs.trezor.lib.protobufB\x12TrezorMessage\
+    Debug\x80\xa6\x1d\x01\
 ";
 
 /// `FileDescriptorProto` object which was a source for this generated file
@@ -3539,10 +3616,11 @@ pub fn file_descriptor() -> &'static ::protobuf::reflect::FileDescriptor {
             messages.push(DebugLinkEraseSdCard::generated_message_descriptor_data());
             messages.push(DebugLinkWatchLayout::generated_message_descriptor_data());
             messages.push(DebugLinkResetDebugEvents::generated_message_descriptor_data());
-            let mut enums = ::std::vec::Vec::with_capacity(3);
+            let mut enums = ::std::vec::Vec::with_capacity(4);
             enums.push(debug_link_decision::DebugSwipeDirection::generated_enum_descriptor_data());
             enums.push(debug_link_decision::DebugButton::generated_enum_descriptor_data());
             enums.push(debug_link_decision::DebugPhysicalButton::generated_enum_descriptor_data());
+            enums.push(debug_link_get_state::DebugWaitType::generated_enum_descriptor_data());
             ::protobuf::reflect::GeneratedFileDescriptor::new_generated(
                 file_descriptor_proto(),
                 deps,

tests/device_handler.py

@@ -2,6 +2,7 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any, Callable
 
 from trezorlib.client import PASSPHRASE_ON_DEVICE
+from trezorlib.messages import DebugWaitType
 from trezorlib.transport import udp
 
 if TYPE_CHECKING:
@@ -42,6 +43,7 @@ class BackgroundDeviceHandler:
         self.client = client
         self.client.ui = NullUI  # type: ignore [NullUI is OK UI]
         self.client.watch_layout(True)
+        self.client.debug.input_wait_type = DebugWaitType.CURRENT_LAYOUT
 
     def run(self, function: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
         """Runs some function that interacts with a device.
@@ -50,8 +52,14 @@ class BackgroundDeviceHandler:
         """
         if self.task is not None:
             raise RuntimeError("Wait for previous task first")
-        self.task = self._pool.submit(function, self.client, *args, **kwargs)
+
-        self.debuglink().wait_layout(wait_for_external_change=True)
+        # make sure we start the wait while a layout is up
+        # TODO should this be part of "wait_for_layout_change"?
+        self.debuglink().read_layout()
+        # from the displayed layout, wait for the first UI change triggered by the
+        # task running in the background
+        with self.debuglink().wait_for_layout_change():
+            self.task = self._pool.submit(function, self.client, *args, **kwargs)
 
     def kill_task(self) -> None:
         if self.task is not None:

tests/device_tests/bitcoin/test_signmessage.py

@@ -380,8 +380,8 @@ def test_signmessage_pagination_trailing_newline(client: Client):
             [
                 # expect address confirmation
                 message_filters.ButtonRequest(code=messages.ButtonRequestType.Other),
-                # expect a ButtonRequest that does not have pagination set
+                # expect a ButtonRequest for a single-page screen
-                message_filters.ButtonRequest(pages=None),
+                message_filters.ButtonRequest(pages=1),
                 messages.MessageSignature,
             ]
         )