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,6 +222,7 @@ message DebugLinkWatchLayout {
  * @next Success
  */
 message DebugLinkResetDebugEvents {
+    option deprecated = true;
 }
 
 

core/src/all_modules.py

@@ -107,6 +107,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
@@ -183,8 +185,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.recovery
 import trezor.ui.layouts.tr.recovery
 trezor.ui.layouts.tr.reset
@@ -193,8 +193,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.recovery
 import trezor.ui.layouts.tt.recovery
 trezor.ui.layouts.tt.reset

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

@@ -20,7 +20,6 @@ from ..keychain import address_n_to_name
 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

@@ -122,10 +122,6 @@ class Progress:
         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

@@ -7,7 +7,7 @@ from . import backup_types
 
 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,16 +8,15 @@ 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,
@@ -25,33 +24,21 @@ if __debug__:
             DebugLinkOptigaSetSecMax,
             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
@@ -62,165 +49,236 @@ 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() -> Awaitable[None]:  # type: ignore [awaitable-is-generator]
-        event_id: int | None, msg: DebugLinkDecision
+        while ui.CURRENT_LAYOUT is None:
+            yield
+
+    async def return_layout_change(
+        ctx: wire.context.Context, detect_deadlock: bool = False
     ) -> None:
+        # set up the wait
+        storage.layout_watcher = True
+
+        # wait for layout change
+        while True:
+            if not detect_deadlock or not layout_change_chan.is_empty():
+                # short-circuit if there is a result already waiting
+                next_layout = await layout_change_chan
+            else:
+                next_layout = await loop.race(
+                    layout_change_chan, _DEADLOCK_DETECT_SLEEP
+                )
+
+            if next_layout is None:
+                # layout close event. loop again
+                continue
+
+            if isinstance(next_layout, ui.Layout):
+                break
+
+            if isinstance(next_layout, int):
+                # sleep result from the deadlock detector
+                raise wire.FirmwareError("layout deadlock detected")
+
+            raise RuntimeError(
+                f"Unexpected layout change: {next_layout}, {type(next_layout)}"
+            )
+
+        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(x: int, y: int, hold_ms: int = 0) -> None:
+        assert isinstance(ui.CURRENT_LAYOUT, ui.Layout)
+        msg = ui.CURRENT_LAYOUT.layout.touch_event(io.TOUCH_START, x, y)
+        ui.CURRENT_LAYOUT._emit_message(msg)
+        ui.CURRENT_LAYOUT._paint()
+
+        if hold_ms:
+            await loop.sleep(hold_ms)
+            workflow.idle_timer.touch()
+
+        if not isinstance(ui.CURRENT_LAYOUT, ui.Layout):
+            return
+        msg = ui.CURRENT_LAYOUT.layout.touch_event(io.TOUCH_END, x, y)
+        ui.CURRENT_LAYOUT._emit_message(msg)
+        ui.CURRENT_LAYOUT._paint()
+
+    async def _layout_press_button(
+        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)
+
+        assert isinstance(ui.CURRENT_LAYOUT, ui.Layout)
+        for btn in buttons:
+            msg = ui.CURRENT_LAYOUT.layout.button_event(io.BUTTON_PRESSED, btn)
+            ui.CURRENT_LAYOUT._emit_message(msg)
+            ui.CURRENT_LAYOUT._paint()
+
+        if hold_ms:
+            await loop.sleep(hold_ms)
+            workflow.idle_timer.touch()
+
+        if not isinstance(ui.CURRENT_LAYOUT, ui.Layout):
+            return
+        for btn in buttons:
+            msg = ui.CURRENT_LAYOUT.layout.button_event(io.BUTTON_RELEASED, btn)
+            ui.CURRENT_LAYOUT._emit_message(msg)
+            ui.CURRENT_LAYOUT._paint()
+
+    if utils.USE_TOUCH:
+
+        async def _layout_swipe(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]
+
+            assert isinstance(ui.CURRENT_LAYOUT, ui.Layout)
+            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 = ui.CURRENT_LAYOUT.layout.touch_event(event, x, y)
+                ui.CURRENT_LAYOUT._emit_message(msg)
+                ui.CURRENT_LAYOUT._paint()
+
+    elif utils.USE_BUTTON:
+
+        def _layout_swipe(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(button)
+
+    else:
+        raise RuntimeError  # No way to swipe with no buttons and no touches
+
+    async def _layout_event(button: DebugButton) -> None:
         from trezor.enums import DebugButton
 
-        if msg.button is not None:
+        assert isinstance(ui.CURRENT_LAYOUT, ui.Layout)
-            if msg.button == DebugButton.NO:
+        if button == DebugButton.NO:
-                await result_chan.put((event_id, trezorui2.CANCELLED))
+            ui.CURRENT_LAYOUT._emit_message(trezorui2.CANCELLED)
-            elif msg.button == DebugButton.YES:
+        elif button == DebugButton.YES:
-                await result_chan.put((event_id, trezorui2.CONFIRMED))
+            ui.CURRENT_LAYOUT._emit_message(trezorui2.CONFIRMED)
-            elif msg.button == DebugButton.INFO:
+        elif button == DebugButton.INFO:
-                await result_chan.put((event_id, trezorui2.INFO))
+            ui.CURRENT_LAYOUT._emit_message(trezorui2.INFO)
-            else:
-                raise RuntimeError(f"Invalid msg.button - {msg.button}")
-        elif msg.input is not None:
-            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 DebugButton")
-            raise RuntimeError("Invalid DebugLinkDecision message")
 
-    async def debuglink_decision_dispatcher() -> None:
+    async def dispatch_DebugLinkDecision(
-        while True:
+        msg: DebugLinkDecision,
-            event_id, msg = await debuglink_decision_chan.take()
+    ) -> DebugLinkState | None:
-            await _dispatch_debuglink_decision(event_id, msg)
+        from trezor import ui, workflow
-
-    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
-                pass
-                # TODO: find out why this sometimes happens on TR when running tests with
-                # "physical" emulator (./emu.py)
-                # raise RuntimeError(
-                #     f"Waiting for event that already happened - {event_id} > {awaited_event_id}"
-                # )
-
-        if awaited_event_id is not None:
-            # Updating last result
-            debug_events.last_result = awaited_event_id
-
-        return content
-
-    async def return_layout_change() -> None:
-        content_tokens = await get_layout_change_content()
-
-        assert isinstance(DEBUG_CONTEXT, context.Context)
-        if storage.layout_watcher is LAYOUT_WATCHER_LAYOUT:
-            await DEBUG_CONTEXT.write(DebugLinkLayout(tokens=content_tokens))
-        else:
-            from trezor.messages import DebugLinkState
-
-            await DEBUG_CONTEXT.write(DebugLinkState(tokens=content_tokens))
-        storage.layout_watcher = LAYOUT_WATCHER_NONE
-
-    async def dispatch_DebugLinkWatchLayout(msg: DebugLinkWatchLayout) -> Success:
-        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
+        assert isinstance(ui.CURRENT_LAYOUT, ui.Layout)
+        layout_change_chan.clear()
 
-        # Touchscreen devices click on specific coordinates, with possible hold
+        try:
-        if (
+            # click on specific coordinates, with possible hold
-            x is not None
+            if x is not None and y is not None:
-            and y is not None
+                await _layout_click(x, y, msg.hold_ms or 0)
-            and utils.INTERNAL_MODEL in ("T2T1", "T3T1", "D001")
+            # press specific button
-        ):
+            elif msg.physical_button is not None:
-            click_chan.publish((debug_events.last_event, x, y, msg.hold_ms))
+                await _layout_press_button(msg.physical_button, msg.hold_ms or 0)
-        # Button devices press specific button
+            elif msg.swipe is not None:
-        elif msg.physical_button is not None and utils.INTERNAL_MODEL in ("T2B1",):
+                await _layout_swipe(msg.swipe)
-            button_chan.publish(
+            elif msg.button is not None:
-                (debug_events.last_event, msg.physical_button, msg.hold_ms)
+                await _layout_event(msg.button)
-            )
+            elif msg.input is not None:
-        else:
+                ui.CURRENT_LAYOUT._emit_message(msg.input)
-            # Will get picked up by _dispatch_debuglink_decision eventually
+            else:
-            debuglink_decision_chan.publish((debug_events.last_event, msg))
+                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())
+        return DebugLinkState(
-            return None
+            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.
@@ -228,6 +286,10 @@ if __debug__:
             REFRESH_INDEX = msg.refresh_index
             storage.save_screen_directory = msg.target_directory
             storage.save_screen = True
+
+            # save the initial screenshot (typically homescreen)
+            screenshot()
+
         else:
             storage.save_screen = False
             display.clear_save()  # clear C buffers
@@ -263,26 +325,6 @@ if __debug__:
             sdcard.power_off()
         return Success()
 
-    def boot() -> None:
-        register = workflow_handlers.register  # local_cache_attribute
-
-        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"]
-        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
-        )
-        register(
-            MessageType.DebugLinkOptigaSetSecMax, dispatch_DebugLinkOptigaSetSecMax
-        )
-
-        loop.schedule(debuglink_decision_dispatcher())
-        if storage.layout_watcher is not LAYOUT_WATCHER_NONE:
-            loop.schedule(return_layout_change())
-
     async def dispatch_DebugLinkOptigaSetSecMax(
         msg: DebugLinkOptigaSetSecMax,
     ) -> Success:
@@ -293,3 +335,89 @@ if __debug__:
             return Success()
         else:
             raise wire.UnexpectedMessage("Optiga not supported")
+
+    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.DebugLinkOptigaSetSecMax: dispatch_DebugLinkOptigaSetSecMax,
+        MessageType.DebugLinkWatchLayout: _no_op,
+        MessageType.DebugLinkResetDebugEvents: _no_op,
+    }
+
+    def boot() -> None:
+        import usb
+
+        loop.schedule(handle_session(usb.iface_debug))

core/src/apps/homescreen/__init__.py

@@ -14,7 +14,7 @@ from apps.common.authorization import is_set_any_session
 async def busyscreen() -> None:
     obj = Busyscreen(busy_expiry_ms())
     try:
-        await obj
+        await obj.get_result()
     finally:
         obj.__del__()
 
@@ -53,7 +53,7 @@ async def homescreen() -> None:
         hold_to_lock=config.has_pin(),
     )
     try:
-        await obj
+        await obj.get_result()
     finally:
         obj.__del__()
 
@@ -72,7 +72,7 @@ async def _lockscreen(screensaver: bool = False) -> None:
             coinjoin_authorized=is_set_any_session(MessageType.AuthorizeCoinJoin),
         )
         try:
-            await obj
+            await obj.get_result()
         finally:
             obj.__del__()
     # Otherwise proceed directly to unlock() call. If the device is already unlocked,

core/src/apps/management/change_language.py

@@ -8,7 +8,7 @@ if TYPE_CHECKING:
     from typing import Callable
 
     from trezor.messages import ChangeLanguage, Success
-    from trezor.ui.layouts.common import ProgressLayout
+    from trezor.ui import ProgressLayout
 
 _CHUNK_SIZE = const(1024)
 

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

@@ -12,20 +12,18 @@ from trezor.ui.layouts.recovery import (  # noqa: F401
 from apps.common import backup_types
 
 if TYPE_CHECKING:
-    from typing import Callable
-
     from trezor.enums import BackupType
+    from trezor.ui.layouts.common import InfoFunc
 
 
 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)
+    send_button_request = True
 
     # Allowing to go back to previous words, therefore cannot use just loop over range(word_count)
     words: list[str] = [""] * word_count
@@ -40,8 +38,10 @@ async def request_mnemonic(
             i,
             word_count,
             is_slip39=backup_types.is_slip39_word_count(word_count),
+            send_button_request=send_button_request,
             prefill_word=words[i],
         )
+        send_button_request = False
 
         # User has decided to go back
         if not word:
@@ -122,7 +122,7 @@ async def homescreen_dialog(
     button_label: str,
     text: str,
     subtext: str | None = None,
-    info_func: Callable | None = None,
+    info_func: InfoFunc | None = None,
     show_info: bool = False,
 ) -> None:
     import storage.recovery as storage_recovery

core/src/apps/misc/get_firmware_hash.py

@@ -2,9 +2,6 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from trezor.messages import FirmwareHash, GetFirmwareHash
-    from trezor.ui.layouts.common import ProgressLayout
-
-_progress_obj: ProgressLayout | None = None
 
 
 async def get_firmware_hash(msg: GetFirmwareHash) -> FirmwareHash:
@@ -14,20 +11,14 @@ async def get_firmware_hash(msg: GetFirmwareHash) -> FirmwareHash:
     from trezor.utils import firmware_hash
 
     workflow.close_others()
-    global _progress_obj
+    progress_obj = progress()
-    _progress_obj = progress()
+
+    def report(progress: int, total: int) -> None:
+        progress_obj.report(1000 * progress // total)
 
     try:
-        hash = firmware_hash(msg.challenge, _render_progress)
+        hash = firmware_hash(msg.challenge, report)
     except ValueError as e:
         raise wire.DataError(str(e))
-    finally:
-        _progress_obj = None
 
     return FirmwareHash(hash=hash)
-
-
-def _render_progress(progress: int, total: int) -> None:
-    global _progress_obj
-    if _progress_obj is not None:
-        _progress_obj.report(1000 * progress // total)

core/src/apps/webauthn/fido2.py

@@ -8,7 +8,8 @@ import storage.device as storage_device
 from trezor import TR, 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,
             TR.fido__device_already_registered,
             TR.fido__already_registered,
             timeout_ms=_POPUP_TIMEOUT_MS,
         )
     else:
-        await show_error_popup(
+        await _show_error_popup(
             title,
             TR.fido__device_not_registered,
             TR.fido__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(
             TR.fido__title_register,
             TR.fido__device_already_registered_with_template,
             TR.fido__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(
             TR.fido__title_verify_user,
             TR.fido__please_enable_pin_protection,
             TR.fido__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(
             TR.fido__title_authenticate,
             TR.fido__not_registered_with_template,
             TR.fido__not_registered,
@@ -1059,6 +1080,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

@@ -65,7 +65,7 @@ async def bootscreen() -> None:
                 lockscreen = Lockscreen(
                     label=storage.device.get_label(), bootscreen=True
                 )
-                await lockscreen
+                await lockscreen.get_result()
                 lockscreen.__del__()
                 await verify_user_pin()
                 storage.init_unlocked()

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

@@ -509,6 +509,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

@@ -669,24 +669,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

@@ -40,6 +40,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
@@ -2761,7 +2762,6 @@ if TYPE_CHECKING:
         input: "str | None"
         x: "int | None"
         y: "int | None"
-        wait: "bool | None"
         hold_ms: "int | None"
         physical_button: "DebugPhysicalButton | None"
 
@@ -2773,7 +2773,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:
@@ -2783,20 +2782,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"
 
@@ -2828,16 +2813,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
 
@@ -2983,26 +2964,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 DebugLinkOptigaSetSecMax(protobuf.MessageType):
 
         @classmethod

core/src/trezor/pin.py

@@ -5,7 +5,7 @@ from . import config
 if TYPE_CHECKING:
     from typing import Any, Container
 
-    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,21 +1,23 @@
 # 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
 from trezorui2 import AttachType, BacklightLevels
 
 if TYPE_CHECKING:
-    from typing import Generic, TypeVar
+    from typing import Any, Callable, Generator, Generic, Iterator, TypeVar
 
-    from trezorui2 import UiResult  # noqa: F401
+    from trezorui2 import LayoutObj, UiResult  # noqa: F401
 
-    T = TypeVar("T")
+    T = TypeVar("T", covariant=True)
 
 else:
-    Generic = [object]
     T = 0
+    Generic = {T: object}
+
 
 # all rendering is done through a singleton of `Display`
 display = Display()
@@ -28,15 +30,14 @@ 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
 
-# storing last transition type, so that next layout can continue nicely
-LAST_TRANSITION_OUT: AttachType | None = None
-
 # in debug mode, display an indicator in top right corner
 if __debug__:
 
@@ -100,102 +101,304 @@ def backlight_fade(val: int, delay: int = 14000, step: int = 15) -> None:
         display.backlight(val)
 
 
-class Result(Exception):
+class Shutdown(Exception):
+    pass
+
+
+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.
     """
-    When components want to trigger layout completion, they do so through
+    global CURRENT_LAYOUT
-    raising an instance of `Result`.
 
-    See `Layout.__iter__` for details.
+    # all transitions must be to/from None
-    """
+    assert (CURRENT_LAYOUT is None) == (layout is not None)
 
-    def __init__(self, value: Any) -> None:
+    CURRENT_LAYOUT = layout
-        super().__init__()
-        self.value = value
 
+    if __debug__ and not isinstance(layout, ProgressLayout):
+        from apps.debug import notify_layout_change
 
-class Cancelled(Exception):
+        notify_layout_change(layout)
-    """
-    Layouts can be explicitly cancelled.  This usually happens when another
-    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.
-    """
 
 
 class Layout(Generic[T]):
-    """
+    """Python-side handler and runner for the Rust based layouts.
-    Abstract class.
 
-    Layouts are top-level components.  Only one layout can be running at the
+    Wrap a `LayoutObj` instance in `Layout` to be able to display the layout, run its
-    same time.  Layouts provide asynchronous interface, so a running task can
+    event loop, and take part in global layout management. See
-    wait for the layout to complete.  Layouts complete when a `Result` is
+    [docs/core/misc/layout-lifecycle.md] for details.
-    raised, usually from some of the child components.
     """
 
-    def finalize(self) -> None:
+    BACKLIGHT_LEVEL = BacklightLevels.NORMAL
-        """
-        Called when the layout is done. Usually overridden to allow cleanup or storing context.
-        """
-        pass
-
-    async def __iter__(self) -> T:
-        """
-        Run the layout and wait until it completes.  Returns the result value.
-        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
-        try:
-            # If any other layout is running (waiting on the layout channel),
-            # we close it with the Cancelled exception, and wait until it is
-            # closed, just to be sure.
-            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
-        finally:
-            self.finalize()
-        return value
-
-    if TYPE_CHECKING:
-
-        def __await__(self) -> Generator[Any, Any, T]:
-            return self.__iter__()  # type: ignore [Coroutine[Any, Any, T@Layout]" is incompatible with "Generator[Any, Any, T@Layout]"]
-
-    else:
-        __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."""
-        raise NotImplementedError
 
     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)
+
+            layout.trace(callback)
+            return "".join(tokens)
+
+        def __str__(self) -> str:
+            return f"{repr(self)}({self._trace(self.layout)[:150]})"
+
+    def __init__(self, layout: LayoutObj[T]) -> None:
+        """Set up a layout."""
+        self.layout = layout
+        self.tasks: set[loop.Task] = set()
+        self.timers: dict[int, loop.Task] = {}
+        self.result_box = loop.mailbox()
+        self.transition_out: AttachType | None = None
+
+    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()
+
+    def is_running(self) -> bool:
+        """True if the layout is in RUNNING state."""
+        return CURRENT_LAYOUT is self
+
+    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.
+
+        If the layout is already RUNNING, do nothing. If the layout is STOPPED, fail.
+        """
+        global CURRENT_LAYOUT
+
+        # 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()
+
+        transition_in = None
+
+        # 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:
+            prev_layout = CURRENT_LAYOUT
+            prev_layout.stop()
+            transition_in = prev_layout.transition_out
+
+        assert CURRENT_LAYOUT is None
+        set_current_layout(self)
+
+        # attach a timer callback and paint self
+        self.layout.attach_timer_fn(self._set_timer, transition_in)
+        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()
+
+        self.transition_out = self.layout.get_transition_out()
+
+        # 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(BacklightLevels.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:
+            return await self.result_box
+        finally:
+            self.stop()
+
+    def request_complete_repaint(self) -> None:
+        """Request a complete repaint of the layout."""
+        msg = self.layout.request_complete_repaint()
+        assert msg is None
+
+    def _paint(self) -> None:
+        """Paint the layout and ensure that homescreen cache is properly invalidated."""
+        import storage.cache as storage_cache
+
+        painted = self.layout.paint()
+        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 _set_timer(self, token: int, deadline: int) -> None:
+        """Timer callback for Rust layouts."""
+
+        async def timer_task() -> None:
+            self.timers.pop(token)
+            result = self.layout.timer(token)
+            self._paint()
+            if result is not None:
+                self.result_box.put(result)
+
+        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 __del__(self) -> None:
+        self.layout.__del__()
 
 
-def wait_until_layout_is_running() -> Awaitable[None]:  # type: ignore [awaitable-return-type]
+class ProgressLayout:
-    while not layout_chan.takers:
+    """Progress layout.
-        yield  # type: ignore [awaitable-return-type]
+
+    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
+        self.transition_out = None
+
+    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(BacklightLevels.NONE)
+        refresh()
+
+    def stop(self) -> None:
+        global CURRENT_LAYOUT
+
+        if CURRENT_LAYOUT is self:
+            CURRENT_LAYOUT = None

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

@@ -1,42 +1,60 @@
 from typing import TYPE_CHECKING
 
-from trezor import workflow
+import trezorui2
+from trezor import 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 Awaitable, Protocol, TypeVar
+    from typing import Awaitable, Callable, TypeVar
 
-    T = TypeVar("T")
-
-    LayoutType = Awaitable
     PropertyType = tuple[str | None, str | bytes | None]
     ExceptionType = BaseException | type[BaseException]
 
-    class ProgressLayout(Protocol):
+    InfoFunc = Callable[[], Awaitable[None]]
-        def report(self, value: int, description: str | None = None) -> None: ...
+
+    T = TypeVar("T")
 
 
-async def button_request(
+async def _button_request(
     br_name: str,
     code: ButtonRequestType = ButtonRequestType.Other,
-    pages: int | None = None,
+    pages: int = 0,
 ) -> None:
     workflow.close_others()
     await context.maybe_call(
-        ButtonRequest(code=code, pages=pages, name=br_name), ButtonAck
+        ButtonRequest(code=code, pages=pages or None, name=br_name), ButtonAck
     )
 
 
 async def interact(
-    layout: LayoutType[T],
+    layout_obj: ui.LayoutObj[T],
-    br_name: str,
+    br_name: str | None,
     br_code: ButtonRequestType = ButtonRequestType.Other,
+    raise_on_cancel: ExceptionType | None = ActionCancelled,
 ) -> T:
-    pages = None
+    # shut down other workflows to prevent them from interfering with the current one
-    if hasattr(layout, "page_count") and layout.page_count() > 1:  # type: ignore [Cannot access attribute "page_count" for class "LayoutType"]
+    workflow.close_others()
-        # We know for certain how many pages the layout will have
+    # start the layout
-        pages = layout.page_count()  # type: ignore [Cannot access attribute "page_count" for class "LayoutType"]
+    layout = ui.Layout(layout_obj)
-    await button_request(br_name, br_code, pages)
+    layout.start()
-    return await layout
+    # send the button request
+    if br_name is not None:
+        await _button_request(br_name, 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
+
+
+def raise_if_not_confirmed(
+    layout_obj: ui.LayoutObj[ui.UiResult],
+    br_name: str | None,
+    br_code: ButtonRequestType = ButtonRequestType.Other,
+    exc: ExceptionType = ActionCancelled,
+) -> Awaitable[None]:
+    action = interact(layout_obj, br_name, br_code, exc)
+    return action  # type: ignore [Type cannot be assigned to type "None"]

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

@@ -1,8 +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 TR, ui
-    from .tr.homescreen import *  # noqa: F401,F403
+
-elif utils.UI_LAYOUT == "MERCURY":
+if TYPE_CHECKING:
-    from .mercury.homescreen import *  # noqa: F401,F403
+    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 notification == TR.homescreen__title_coinjoin_authorized:
+                level = 3
+            elif notification == TR.homescreen__title_experimental_mode:
+                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.BacklightLevels.LOW
+
+    def __init__(
+        self,
+        label: str | None,
+        bootscreen: bool = False,
+        coinjoin_authorized: bool = False,
+    ) -> None:
+        self.bootscreen = bootscreen
+        if bootscreen:
+            self.BACKLIGHT_LEVEL = ui.BacklightLevels.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=TR.coinjoin__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/mercury/recovery.py

@@ -71,30 +71,24 @@ async def show_remaining_shares(
             pages.append((title, words))
 
     await raise_if_not_confirmed(
-        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:
     await raise_if_not_confirmed(
-        interact(
+        trezorui2.show_group_share_success(
-            RustLayout(
+            lines=[
-                trezorui2.show_group_share_success(
+                TR.recovery__you_have_entered,
-                    lines=[
+                TR.recovery__share_num_template.format(share_index + 1),
-                        TR.recovery__you_have_entered,
+                TR.words__from,
-                        TR.recovery__share_num_template.format(share_index + 1),
+                TR.recovery__group_num_template.format(group_index + 1),
-                        TR.words__from,
+            ],
-                        TR.recovery__group_num_template.format(group_index + 1),
+        ),
-                    ],
+        "share_success",
-                )
+        ButtonRequestType.Other,
-            ),
-            "share_success",
-            ButtonRequestType.Other,
-        )
     )
 
 

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

@@ -106,13 +106,11 @@ async def slip39_show_checklist(
 ) -> None:
     items = _slip_39_checklist_items(step, advanced, count, threshold)
     result = await interact(
-        RustLayout(
+        trezorui2.show_checklist(
-            trezorui2.show_checklist(
+            title=TR.reset__title_shamir_backup,
-                title=TR.reset__title_shamir_backup,
+            button=TR.buttons__continue,
-                button=TR.buttons__continue,
+            active=step,
-                active=step,
+            items=items,
-                items=items,
-            )
         ),
         "slip39_checklist",
         ButtonRequestType.ResetDevice,
@@ -301,25 +299,21 @@ async def show_intro_backup(single_share: bool, num_of_words: int | None) -> Non
         description = TR.backup__info_multi_share_backup
 
     await interact(
-        RustLayout(
+        trezorui2.show_info(
-            trezorui2.show_info(
+            title=TR.backup__title_create_wallet_backup, description=description
-                title=TR.backup__title_create_wallet_backup, description=description
-            )
         ),
-        "backup_warning",
+        "backup_intro",
         ButtonRequestType.ResetDevice,
     )
 
 
 async def show_warning_backup() -> None:
     result = await interact(
-        RustLayout(
+        trezorui2.show_warning(
-            trezorui2.show_warning(
+            title=TR.words__important,
-                title=TR.words__important,
+            value=TR.reset__never_make_digital_copy,
-                value=TR.reset__never_make_digital_copy,
+            button="",
-                button="",
+            allow_cancel=False,
-                allow_cancel=False,
-            )
         ),
         "backup_warning",
         ButtonRequestType.ResetDevice,
@@ -342,20 +336,16 @@ async def show_reset_warning(
     button: str | None = None,
     br_code: ButtonRequestType = ButtonRequestType.Warning,
 ) -> None:
-    await raise_if_not_confirmed(
+    await interact(
-        interact(
+        trezorui2.show_warning(
-            RustLayout(
+            title=subheader or "",
-                trezorui2.show_warning(
+            description=content,
-                    title=subheader or "",
+            value="",
-                    description=content,
+            button="",
-                    value="",
+            allow_cancel=False,
-                    button="",
+        ),
-                    allow_cancel=False,
+        br_name,
-                )
+        br_code,
-            ),
-            br_name,
-            br_code,
-        )
     )
 
 

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

@@ -1,13 +1,6 @@
-from typing import TYPE_CHECKING
-
 import trezorui2
 from trezor import TR, config, ui, utils
 
-if TYPE_CHECKING:
-    from typing import Any
-
-    from .common import ProgressLayout
-
 
 def _storage_message_to_str(message: config.StorageMessage | None) -> str | None:
     from trezor import TR
@@ -28,37 +21,22 @@ def _storage_message_to_str(message: config.StorageMessage | None) -> str | None
     raise RuntimeError  # unknown message
 
 
-class RustProgress:
-    def __init__(
-        self,
-        layout: Any,
-    ):
-        self.layout = layout
-        ui.backlight_fade(ui.BacklightLevels.DIM)
-        self.layout.attach_timer_fn(self.set_timer, None)
-        if self.layout.paint():
-            ui.refresh()
-        ui.backlight_fade(ui.BacklightLevels.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
-        if self.layout.paint():
-            ui.refresh()
-
-
 def progress(
     description: str | None = None,
     title: str | None = None,
     indeterminate: bool = False,
-) -> ProgressLayout:
+) -> ui.ProgressLayout:
     if description is None:
         description = TR.progress__please_wait  # def_arg
 
-    return RustProgress(
+    if title is not None:
+        title = title.upper()
+    elif utils.MODEL != "T2B1":
+        # on TT, uppercase the description which ends up on top of the screen
+        # when no title is set
+        description = description.upper()
+
+    return ui.ProgressLayout(
         layout=trezorui2.show_progress(
             description=description,
             title=title,
@@ -67,27 +45,27 @@ def progress(
     )
 
 
-def bitcoin_progress(message: str) -> ProgressLayout:
+def bitcoin_progress(message: str) -> ui.ProgressLayout:
     return progress(message)
 
 
-def coinjoin_progress(message: str) -> ProgressLayout:
+def coinjoin_progress(message: str) -> ui.ProgressLayout:
-    return RustProgress(
+    return ui.ProgressLayout(
         layout=trezorui2.show_progress_coinjoin(title=message, indeterminate=False)
     )
 
 
-def pin_progress(title: config.StorageMessage, description: str) -> ProgressLayout:
+def pin_progress(title: config.StorageMessage, description: str) -> ui.ProgressLayout:
     return progress(description=description, title=_storage_message_to_str(title))
 
 
 if not utils.BITCOIN_ONLY:
 
-    def monero_keyimage_sync_progress() -> ProgressLayout:
+    def monero_keyimage_sync_progress() -> ui.ProgressLayout:
         return progress(TR.progress__syncing)
 
-    def monero_live_refresh_progress() -> ProgressLayout:
+    def monero_live_refresh_progress() -> ui.ProgressLayout:
         return progress(TR.progress__refreshing, indeterminate=True)
 
-    def monero_transaction_progress_inner() -> ProgressLayout:
+    def monero_transaction_progress_inner() -> ui.ProgressLayout:
         return progress(TR.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,
-            title=header,
+        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,22 +27,17 @@ 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:
     from trezor import TR
 
-    confirm = RustLayout(
+    confirm = trezorui2.confirm_action(
-        trezorui2.confirm_action(
+        title=TR.fido__title_reset,
-            title=TR.fido__title_reset,
+        description=TR.fido__wanna_erase_credentials,
-            description=TR.fido__wanna_erase_credentials,
+        action=None,
-            action=None,
+        verb_cancel="",
-            verb_cancel="",
+        verb=TR.buttons__confirm,
-            verb=TR.buttons__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,128 +0,0 @@
-from typing import TYPE_CHECKING
-
-import storage.cache as storage_cache
-import trezorui2
-from trezor import TR, 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:
-        if 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:
-            if notification == TR.homescreen__title_experimental_mode:
-                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)
-            if 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:
-        from trezor import TR
-
-        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        super().__init__(
-            layout=trezorui2.show_progress_coinjoin(
-                title=TR.coinjoin__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/recovery.py

@@ -1,17 +1,20 @@
-from typing import Callable, Iterable
+from typing import TYPE_CHECKING, Awaitable, Iterable
 
 import trezorui2
-from trezor import TR
+from trezor import TR, ui
 from trezor.enums import ButtonRequestType, RecoveryType
 
 from ..common import interact
-from . import RustLayout, raise_if_not_confirmed, show_warning
+from . import show_warning
+
+if TYPE_CHECKING:
+    from ..common import InfoFunc
 
 
 async def request_word_count(recovery_type: RecoveryType) -> int:
     count = await interact(
-        RustLayout(trezorui2.select_word_count(recovery_type=recovery_type)),
+        trezorui2.select_word_count(recovery_type=recovery_type),
-        "word_count",
+        "recovery_word_count",
         ButtonRequestType.MnemonicWordCount,
     )
     # It can be returning a string (for example for __debug__ in tests)
@@ -19,26 +22,31 @@ async def request_word_count(recovery_type: RecoveryType) -> int:
 
 
 async def request_word(
-    word_index: int, word_count: int, is_slip39: bool, prefill_word: str = ""
+    word_index: int,
+    word_count: int,
+    is_slip39: bool,
+    send_button_request: bool,
+    prefill_word: str = "",
 ) -> str:
     prompt = TR.recovery__word_x_of_y_template.format(word_index + 1, word_count)
 
     can_go_back = word_index > 0
 
     if is_slip39:
-        word_choice = RustLayout(
+        keyboard = trezorui2.request_slip39(
-            trezorui2.request_slip39(
+            prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
-                prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
-            )
-        )
-    else:
-        word_choice = RustLayout(
-            trezorui2.request_bip39(
-                prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
-            )
         )
 
-    word: str = await word_choice
+    else:
+        keyboard = trezorui2.request_bip39(
+            prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
+        )
+
+    word: str = await interact(
+        keyboard,
+        "mnemonic" if send_button_request else None,
+        ButtonRequestType.MnemonicInput,
+    )
     return word
 
 
@@ -50,22 +58,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=[
-                        TR.recovery__you_have_entered,
+                TR.recovery__you_have_entered,
-                        TR.recovery__share_num_template.format(share_index + 1),
+                TR.recovery__share_num_template.format(share_index + 1),
-                        TR.words__from,
+                TR.words__from,
-                        TR.recovery__group_num_template.format(group_index + 1),
+                TR.recovery__group_num_template.format(group_index + 1),
-                    ],
+            ],
-                )
+        ),
-            ),
+        "share_success",
-            "share_success",
+        ButtonRequestType.Other,
-            ButtonRequestType.Other,
-        )
     )
 
 
@@ -97,7 +103,7 @@ async def continue_recovery(
     button_label: str,
     text: str,
     subtext: str | None,
-    info_func: Callable | None,
+    info_func: InfoFunc | None,
     recovery_type: RecoveryType,
     show_info: bool = False,
 ) -> bool:
@@ -114,40 +120,29 @@ async def continue_recovery(
     if subtext:
         text += f"\n\n{subtext}"
 
-    while True:
+    homepage = trezorui2.confirm_recovery(
-        homepage = RustLayout(
+        title="",
-            trezorui2.confirm_recovery(
+        description=text,
-                title="",
+        button=button_label.upper(),
-                description=text,
+        recovery_type=recovery_type,
-                button=button_label,
+        info_button=False,
-                recovery_type=recovery_type,
+        show_info=show_info,  # type: ignore [No parameter named "show_info"]
-                info_button=False,
+    )
-                show_info=show_info,  # type: ignore [No parameter named "show_info"]
+    result = await interact(
-            )
+        homepage,
-        )
+        "recovery",
-        result = await interact(
+        ButtonRequestType.RecoveryHomepage,
-            homepage,
+        raise_on_cancel=None,
-            "recovery",
+    )
-            ButtonRequestType.RecoveryHomepage,
+    return result is trezorui2.CONFIRMED
-        )
-        if result is trezorui2.CONFIRMED:
-            return True
-
-        # user has chosen to abort, confirm the choice
-        try:
-            await _confirm_abort(recovery_type != RecoveryType.NormalRecovery)
-        except ActionCancelled:
-            pass
-        else:
-            return False
 
 
-async def show_recovery_warning(
+def show_recovery_warning(
     br_name: str,
     content: str,
     subheader: str | None = None,
     button: str | None = None,
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[ui.UiResult]:
     button = button or TR.buttons__try_again  # def_arg
-    await show_warning(br_name, content, subheader, button, br_code)
+    return show_warning(br_name, content, subheader, button, br_code)

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

@@ -3,13 +3,17 @@ from typing import Sequence
 import trezorui2
 from trezor import TR
 from trezor.enums import ButtonRequestType
-from trezor.wire import ActionCancelled
 
-from ..common import interact
+from ..common import interact, raise_if_not_confirmed
-from . import RustLayout, confirm_action, show_success, show_warning
+from . import confirm_action, show_success, show_warning
 
 CONFIRMED = trezorui2.CONFIRMED  # global_import_cache
 
+if TYPE_CHECKING:
+    from typing import Awaitable, Sequence
+
+    from trezor.enums import BackupType
+
 
 async def show_share_words(
     share_words: Sequence[str],
@@ -45,13 +49,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_name,
             br_code,
+            raise_on_cancel=None,
         )
         if result is CONFIRMED:
             break
@@ -82,13 +85,14 @@ async def select_word(
     while len(words) < 3:
         words.append(words[-1])
 
-    word_ordinal = format_ordinal(checked_index + 1)
+    word_ordinal = format_ordinal(checked_index + 1).upper()
-    result = await RustLayout(
+    result = await interact(
         trezorui2.select_word(
             title="",
             description=TR.reset__select_word_template.format(word_ordinal),
             words=(words[0].lower(), words[1].lower(), words[2].lower()),
-        )
+        ),
+        None,
     )
     if __debug__ and isinstance(result, str):
         return result
@@ -96,12 +100,12 @@ async def select_word(
     return words[result]
 
 
-async def slip39_show_checklist(
+def slip39_show_checklist(
     step: int,
     advanced: bool,
     count: int | None = None,
     threshold: int | None = None,
-) -> None:
+) -> Awaitable[None]:
     items = (
         (
             TR.reset__slip39_checklist_num_shares,
@@ -116,20 +120,16 @@ async def slip39_show_checklist(
         )
     )
 
-    result = await interact(
+    return raise_if_not_confirmed(
-        RustLayout(
+        trezorui2.show_checklist(
-            trezorui2.show_checklist(
+            title=TR.reset__slip39_checklist_title,
-                title=TR.reset__slip39_checklist_title,
+            button=TR.buttons__continue,
-                button=TR.buttons__continue,
+            active=step,
-                active=step,
+            items=items,
-                items=items,
-            )
         ),
         "slip39_checklist",
         ButtonRequestType.ResetDevice,
     )
-    if result is not CONFIRMED:
-        raise ActionCancelled
 
 
 async def _prompt_number(
@@ -139,13 +139,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,
-            title=title,
+        count=count,
-            count=count,
+        min_count=min_count,
-            min_count=min_count,
+        max_count=max_count,
-            max_count=max_count,
-        )
     )
 
     result = await interact(
@@ -225,12 +223,12 @@ async def slip39_prompt_number_of_shares(
     )
 
 
-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(
         TR.reset__title_number_of_groups,
         count,
         min_count,
@@ -239,12 +237,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(
         TR.reset__title_group_threshold,
         count,
         min_count,
@@ -253,15 +251,15 @@ async def slip39_advanced_prompt_group_threshold(num_of_groups: int) -> int:
     )
 
 
-async def show_intro_backup(single_share: bool, num_of_words: int | None) -> None:
+def show_intro_backup(single_share: bool, num_of_words: int | None) -> Awaitable[None]:
     if single_share:
         assert num_of_words is not None
         description = TR.backup__info_single_share_backup.format(num_of_words)
     else:
         description = TR.backup__info_multi_share_backup
 
-    await confirm_action(
+    return confirm_action(
-        "backup_warning",
+        "backup_intro",
         title=TR.backup__title_backup_wallet,
         verb=TR.buttons__continue,
         description=description,
@@ -270,8 +268,8 @@ async def show_intro_backup(single_share: bool, num_of_words: int | None) -> Non
     )
 
 
-async def show_warning_backup() -> None:
+def show_warning_backup() -> Awaitable[trezorui2.UiResult]:
-    await show_warning(
+    return show_warning(
         "backup_warning",
         TR.words__title_remember,
         TR.reset__never_make_digital_copy,
@@ -280,8 +278,8 @@ async def show_warning_backup() -> None:
     )
 
 
-async def show_success_backup() -> None:
+def show_success_backup() -> Awaitable[None]:
-    await confirm_action(
+    return confirm_action(
         "success_backup",
         TR.reset__title_backup_is_done,
         description=TR.words__keep_it_safe,
@@ -291,16 +289,16 @@ async def show_success_backup() -> None:
     )
 
 
-async def show_reset_warning(
+def show_reset_warning(
     br_name: str,
     content: str,
     subheader: str | None = None,
     button: str | None = None,
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[trezorui2.UiResult]:
     button = button or TR.buttons__try_again  # def_arg
 
-    await show_warning(
+    return show_warning(
         br_name,
         subheader or "",
         content,

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

@@ -1,50 +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, ...]:
-            return (
-                self.handle_input_and_rendering(),
-                self.handle_timers(),
-                self.handle_swipe(),
-                self.handle_debug_confirm(),
-            )
-
-        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)
-                if self.layout.paint():
-                    ui.refresh()
-                if msg is not None:
-                    raise Result(msg)
-
-    _RustFidoLayout = _RustFidoLayoutImpl
-
-else:
-    _RustFidoLayout = RustLayout
 
 
 async def confirm_fido(
@@ -54,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,
-            title=header,
+        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):
@@ -78,7 +50,7 @@ async def confirm_fido(
 async def confirm_fido_reset() -> bool:
     from trezor import TR
 
-    confirm = RustLayout(
+    confirm = ui.Layout(
         trezorui2.confirm_action(
             title=TR.fido__title_reset,
             action=TR.fido__erase_credentials,
@@ -86,4 +58,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,143 +0,0 @@
-from typing import TYPE_CHECKING
-
-import storage.cache as storage_cache
-import trezorui2
-from trezor import TR, 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:
-        if 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, ...]:
-            return (
-                self.handle_input_and_rendering(),
-                self.handle_timers(),
-                self.handle_click_signal(),  # so we can receive debug events
-            )
-
-
-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:
-            if notification == TR.homescreen__title_coinjoin_authorized:
-                level = 3
-            elif notification == TR.homescreen__title_experimental_mode:
-                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)
-            if 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
-        self.backlight_level = ui.BacklightLevels.LOW
-        if bootscreen:
-            self.backlight_level = ui.BacklightLevels.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:
-        from trezor import TR
-
-        skip = storage_cache.homescreen_shown is self.RENDER_INDICATOR
-        super().__init__(
-            layout=trezorui2.show_progress_coinjoin(
-                title=TR.coinjoin__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/recovery.py

@@ -1,63 +1,58 @@
-from typing import Callable, Iterable
+from typing import TYPE_CHECKING, Awaitable, Iterable
 
 import trezorui2
-from trezor import TR
+from trezor import TR, ui
-from trezor.enums import ButtonRequestType, RecoveryType
+from trezor.enums import ButtonRequestType
 
 from ..common import interact
-from . import RustLayout, raise_if_not_confirmed
 
-CONFIRMED = trezorui2.CONFIRMED  # global_import_cache
+if TYPE_CHECKING:
-INFO = trezorui2.INFO  # global_import_cache
+    from trezor.enums import RecoveryType
 
-
+    from ..common import InfoFunc
-async def _homepage_with_info(
-    dialog: RustLayout,
-    info_func: Callable,
-) -> trezorui2.UiResult:
-    while True:
-        result = await dialog
-
-        if result is INFO:
-            await info_func()
-            dialog.request_complete_repaint()
-        else:
-            return result
 
 
 async def request_word_count(recovery_type: RecoveryType) -> int:
-    selector = RustLayout(trezorui2.select_word_count(recovery_type=recovery_type))
+    count = await interact(
-    count = await interact(selector, "word_count", ButtonRequestType.MnemonicWordCount)
+        trezorui2.select_word_count(recovery_type=recovery_type),
+        "word_count",
+        ButtonRequestType.MnemonicWordCount,
+    )
     return int(count)
 
 
 async def request_word(
-    word_index: int, word_count: int, is_slip39: bool, prefill_word: str = ""
+    word_index: int,
+    word_count: int,
+    is_slip39: bool,
+    send_button_request: bool,
+    prefill_word: str = "",
 ) -> str:
     prompt = TR.recovery__type_word_x_of_y_template.format(word_index + 1, word_count)
     can_go_back = word_index > 0
     if is_slip39:
-        keyboard = RustLayout(
+        keyboard = trezorui2.request_slip39(
-            trezorui2.request_slip39(
+            prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
-                prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
-            )
-        )
-    else:
-        keyboard = RustLayout(
-            trezorui2.request_bip39(
-                prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
-            )
         )
 
-    word: str = await keyboard
+    else:
+        keyboard = trezorui2.request_bip39(
+            prompt=prompt, prefill_word=prefill_word, can_go_back=can_go_back
+        )
+
+    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
 
@@ -83,31 +78,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=[
-                        TR.recovery__you_have_entered,
+                TR.recovery__you_have_entered,
-                        TR.recovery__share_num_template.format(share_index + 1),
+                TR.recovery__share_num_template.format(share_index + 1),
-                        TR.words__from,
+                TR.words__from,
-                        TR.recovery__group_num_template.format(group_index + 1),
+                TR.recovery__group_num_template.format(group_index + 1),
-                    ],
+            ],
-                )
+        ),
-            ),
+        "share_success",
-            "share_success",
+        ButtonRequestType.Other,
-            ButtonRequestType.Other,
-        )
     )
 
 
@@ -139,67 +130,56 @@ async def continue_recovery(
     button_label: str,
     text: str,
     subtext: str | None,
-    info_func: Callable | None,
+    info_func: InfoFunc | None,
     recovery_type: RecoveryType,
     show_info: bool = False,
 ) -> bool:
-    from trezor.wire import ActionCancelled
-
-    from ..common import button_request
-
     if show_info:
         # Show this just one-time
         description = TR.recovery__enter_each_word
     else:
         description = subtext or ""
 
+    homepage = trezorui2.confirm_recovery(
+        title=text,
+        description=description,
+        button=button_label,
+        recovery_type=recovery_type,
+        info_button=info_func is not None,
+    )
+
+    send_button_request = True
     while True:
-        homepage = RustLayout(
+        result = await interact(
-            trezorui2.confirm_recovery(
+            homepage,
-                title=text,
+            "recovery" if send_button_request else None,
-                description=description,
+            ButtonRequestType.RecoveryHomepage,
-                button=button_label,
+            raise_on_cancel=None,
-                recovery_type=recovery_type,
-                info_button=info_func is not None,
-            )
         )
+        send_button_request = False
 
-        await button_request("recovery", ButtonRequestType.RecoveryHomepage)
+        if info_func is not None and result is trezorui2.INFO:
-
+            await info_func()
-        result = (
-            await homepage
-            if info_func is None
-            else await _homepage_with_info(homepage, info_func)
-        )
-        if result is CONFIRMED:
-            return True
-        try:
-            await _confirm_abort(recovery_type != RecoveryType.NormalRecovery)
-        except ActionCancelled:
-            pass
         else:
-            return False
+            return result is trezorui2.CONFIRMED
 
 
-async def show_recovery_warning(
+def show_recovery_warning(
     br_name: str,
     content: str,
     subheader: str | None = None,
     button: str | None = None,
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[ui.UiResult]:
     button = button or TR.buttons__try_again  # def_arg
-    await raise_if_not_confirmed(
+
-        interact(
+    return interact(
-            RustLayout(
+        trezorui2.show_warning(
-                trezorui2.show_warning(
+            title=content,
-                    title=content,
+            description=subheader or "",
-                    description=subheader or "",
+            button=button,
-                    button=button,
+            allow_cancel=False,
-                    allow_cancel=False,
+        ),
-                )
+        br_name,
-            ),
+        br_code,
-            br_name,
-            br_code,
-        )
     )

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

@@ -1,12 +1,11 @@
-from typing import Callable, Sequence
+from typing import Awaitable, Callable, Sequence
 
 import trezorui2
 from trezor import TR
 from trezor.enums import ButtonRequestType
-from trezor.wire import ActionCancelled
 
-from ..common import interact
+from ..common import interact, raise_if_not_confirmed
-from . import RustLayout, raise_if_not_confirmed, show_success
+
 
 CONFIRMED = trezorui2.CONFIRMED  # global_import_cache
 
@@ -35,11 +34,11 @@ def _split_share_into_pages(share_words: Sequence[str], per_page: int = 4) -> li
     return pages
 
 
-async def show_share_words(
+def show_share_words(
     share_words: Sequence[str],
     share_index: int | None = None,
     group_index: int | None = None,
-) -> None:
+) -> Awaitable[None]:
     if share_index is None:
         title = TR.reset__recovery_wallet_backup_title
     elif group_index is None:
@@ -51,18 +50,14 @@ async def show_share_words(
 
     pages = _split_share_into_pages(share_words)
 
-    result = await interact(
+    return raise_if_not_confirmed(
-        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(
@@ -88,14 +83,15 @@ async def select_word(
     while len(words) < 3:
         words.append(words[-1])
 
-    result = await RustLayout(
+    result = await interact(
         trezorui2.select_word(
             title=title,
             description=TR.reset__select_word_x_of_y_template.format(
                 checked_index + 1, count
             ),
             words=(words[0], words[1], words[2]),
-        )
+        ),
+        None,
     )
     if __debug__ and isinstance(result, str):
         return result
@@ -103,12 +99,12 @@ async def select_word(
     return words[result]
 
 
-async def slip39_show_checklist(
+def slip39_show_checklist(
     step: int,
     advanced: bool,
     count: int | None = None,
     threshold: int | None = None,
-) -> None:
+) -> Awaitable[None]:
     items = (
         (
             TR.reset__slip39_checklist_set_num_shares,
@@ -123,20 +119,16 @@ async def slip39_show_checklist(
         )
     )
 
-    result = await interact(
+    return raise_if_not_confirmed(
-        RustLayout(
+        trezorui2.show_checklist(
-            trezorui2.show_checklist(
+            title=TR.reset__slip39_checklist_title,
-                title=TR.reset__slip39_checklist_title,
+            button=TR.buttons__continue,
-                button=TR.buttons__continue,
+            active=step,
-                active=step,
+            items=items,
-                items=items,
-            )
         ),
         "slip39_checklist",
         ButtonRequestType.ResetDevice,
     )
-    if result != CONFIRMED:
-        raise ActionCancelled
 
 
 async def _prompt_number(
@@ -148,14 +140,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,
-            title=title,
+        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:
@@ -163,31 +153,33 @@ 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 RustLayout(
+        await interact(
             trezorui2.show_simple(
                 title=None,
                 description=info(value),
                 button=TR.buttons__ok_i_understand,
-            )
+            ),
+            None,
+            raise_on_cancel=None,
         )
-        num_input.request_complete_repaint()
 
 
-async def slip39_prompt_threshold(
+def slip39_prompt_threshold(
     num_of_shares: int, group_id: int | None = None
-) -> int:
+) -> Awaitable[int]:
     count = num_of_shares // 2 + 1
     # min value of share threshold is 2 unless the number of shares is 1
     # number of shares 1 is possible in advanced slip39
@@ -230,7 +222,7 @@ async def slip39_prompt_threshold(
             text += " " + TR.reset__to_form_group_template.format(group_id + 1)
         return text
 
-    return await _prompt_number(
+    return _prompt_number(
         TR.reset__title_set_threshold,
         description,
         info,
@@ -241,9 +233,7 @@ async def slip39_prompt_threshold(
     )
 
 
-async def slip39_prompt_number_of_shares(
+def slip39_prompt_number_of_shares(group_id: int | None = None) -> Awaitable[int]:
-    num_words: int, group_id: int | None = None
-) -> int:
     count = 5
     min_count = 1
     max_count = 16
@@ -266,7 +256,7 @@ async def slip39_prompt_number_of_shares(
             num_words, group_id + 1
         )
 
-    return await _prompt_number(
+    return _prompt_number(
         TR.reset__title_set_number_of_shares,
         description,
         lambda i: info,
@@ -277,14 +267,14 @@ async def slip39_prompt_number_of_shares(
     )
 
 
-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
     description = TR.reset__group_description
     info = TR.reset__group_info
 
-    return await _prompt_number(
+    return _prompt_number(
         TR.reset__title_set_number_of_groups,
         lambda i: description,
         lambda i: info,
@@ -295,14 +285,14 @@ 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
     description = TR.reset__required_number_of_groups
     info = TR.reset__advanced_group_threshold_info
 
-    return await _prompt_number(
+    return _prompt_number(
         TR.reset__title_set_group_threshold,
         lambda i: description,
         lambda i: info,
@@ -313,44 +303,40 @@ async def slip39_advanced_prompt_group_threshold(num_of_groups: int) -> int:
     )
 
 
-async def show_intro_backup(single_share: bool, num_of_words: int | None) -> None:
+def show_intro_backup(single_share: bool, num_of_words: int | None) -> Awaitable[None]:
     if single_share:
         assert num_of_words is not None
         description = TR.backup__info_single_share_backup.format(num_of_words)
     else:
         description = TR.backup__info_multi_share_backup
 
-    await interact(
+    return raise_if_not_confirmed(
-        RustLayout(
+        trezorui2.show_info(
-            trezorui2.show_info(
+            title="",
-                title="",
+            button=TR.buttons__continue,
-                button=TR.buttons__continue,
+            description=description,
-                description=description,
+            allow_cancel=False,
-                allow_cancel=False,
-            )
         ),
-        "backup_warning",
+        "backup_intro",
         ButtonRequestType.ResetDevice,
     )
 
 
-async def show_warning_backup() -> None:
+def show_warning_backup() -> Awaitable[trezorui2.UiResult]:
-    result = await interact(
+    return interact(
-        RustLayout(
+        trezorui2.show_info(
-            trezorui2.show_info(
+            title=TR.reset__never_make_digital_copy,
-                title=TR.reset__never_make_digital_copy,
+            button=TR.buttons__ok_i_understand,
-                button=TR.buttons__ok_i_understand,
+            allow_cancel=False,
-                allow_cancel=False,
-            )
         ),
         "backup_warning",
         ButtonRequestType.ResetDevice,
     )
-    if result != CONFIRMED:
-        raise ActionCancelled
 
 
 async def show_success_backup() -> None:
+    from . import show_success
+
     await show_success(
         "success_backup",
         TR.reset__use_your_backup,
@@ -358,27 +344,23 @@ async def show_success_backup() -> None:
     )
 
 
-async def show_reset_warning(
+def show_reset_warning(
     br_name: str,
     content: str,
     subheader: str | None = None,
     button: str | None = None,
     br_code: ButtonRequestType = ButtonRequestType.Warning,
-) -> None:
+) -> Awaitable[trezorui2.UiResult]:
     button = button or TR.buttons__try_again  # def_arg
-    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,
+        ),
-                    allow_cancel=False,
+        br_name,
-                )
+        br_code,
-            ),
-            br_name,
-            br_code,
-        )
     )
 
 
@@ -387,6 +369,8 @@ async def show_share_confirmation_success(
     num_of_shares: int | None = None,
     group_index: int | None = None,
 ) -> None:
+    from . import show_success
+
     if share_index is None or num_of_shares is None:
         # it is a BIP39 or a 1-of-1 SLIP39 backup
         subheader = TR.reset__finished_verifying_wallet_backup

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, models, 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__)
@@ -372,6 +383,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
@@ -386,7 +420,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
 
@@ -394,6 +427,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."""
@@ -415,7 +453,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},
@@ -426,13 +469,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)
 
@@ -448,11 +490,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
@@ -463,11 +514,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 is not models.T1B1) and not self.legacy_debug:
@@ -511,56 +589,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:
@@ -579,127 +703,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))
@@ -733,44 +738,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 is models.T1B1 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"
         )
@@ -778,6 +774,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
@@ -835,15 +834,9 @@ class DebugUI:
                 self.debuglink.press_yes()
 
     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)
             self._default_input_flow(br)
         elif self.input_flow is self.INPUT_FLOW_DONE:
             raise AssertionError("input flow ended prematurely")
@@ -855,7 +848,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")
@@ -866,7 +859,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

@@ -552,6 +552,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
@@ -4028,7 +4034,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__(
@@ -4036,7 +4042,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

@@ -510,12 +510,17 @@ def format_message(
         return printable / len(bytes) > 0.8
 
     def pformat(name: str, value: t.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 {
@@ -3560,53 +3633,56 @@ impl ::protobuf::reflect::ProtobufValue for DebugLinkOptigaSetSecMax {
 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\x19DebugLinkResetDebugEvents\"\x1a\n\x18DebugLinkOptigaSetSec\
+    -\n\x13DebugLinkFlashErase\x12\x16\n\x06sector\x18\x01\x20\x01(\rR\x06se\
-    MaxB=\n#com.satoshilabs.trezor.lib.protobufB\x12TrezorMessageDebug\x80\
+    ctor\".\n\x14DebugLinkEraseSdCard\x12\x16\n\x06format\x18\x01\x20\x01(\
-    \xa6\x1d\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\x01\"\x1a\n\x18DebugLinkOptigaSetSecMaxB=\n#com.satoshilabs.t\
+    rezor.lib.protobufB\x12TrezorMessageDebug\x80\xa6\x1d\x01\
 ";
 
 /// `FileDescriptorProto` object which was a source for this generated file
@@ -3644,10 +3720,11 @@ pub fn file_descriptor() -> &'static ::protobuf::reflect::FileDescriptor {
             messages.push(DebugLinkWatchLayout::generated_message_descriptor_data());
             messages.push(DebugLinkResetDebugEvents::generated_message_descriptor_data());
             messages.push(DebugLinkOptigaSetSecMax::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

@@ -392,8 +392,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,
             ]
         )