fix(core): TS3 device freeze

[no changelog]

(cherry picked from commit b560ce53c0)

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

@@ -224,11 +224,17 @@ class RustLayout(LayoutParentType[T]):
     else:
 
         def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
-            return (
-                self.handle_timers(),
-                self.handle_input_and_rendering(),
-                self.handle_usb(context.get_context()),
-            )
+            if context.CURRENT_CONTEXT:
+                return (
+                    self.handle_timers(),
+                    self.handle_input_and_rendering(),
+                    self.handle_usb(context.get_context()),
+                )
+            else:
+                return (
+                    self.handle_timers(),
+                    self.handle_input_and_rendering(),
+                )
 
     def _first_paint(self) -> None:
         self._paint()

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

@@ -21,8 +21,6 @@ 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 = ""
 ) -> str:
-    from trezor.wire.context import wait
-
     prompt = TR.recovery__word_x_of_y_template.format(word_index + 1, word_count)
 
     can_go_back = word_index > 0
@@ -40,7 +38,7 @@ async def request_word(
             )
         )
 
-    word: str = await wait(word_choice)
+    word: str = await word_choice
     return word
 
 

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

@@ -74,7 +74,6 @@ async def select_word(
     group_index: int | None = None,
 ) -> str:
     from trezor.strings import format_ordinal
-    from trezor.wire.context import wait
 
     # It may happen (with a very low probability)
     # that there will be less than three unique words to choose from.
@@ -84,13 +83,11 @@ async def select_word(
         words.append(words[-1])
 
     word_ordinal = format_ordinal(checked_index + 1)
-    result = await wait(
-        RustLayout(
-            trezorui2.select_word(
-                title="",
-                description=TR.reset__select_word_template.format(word_ordinal),
-                words=(words[0].lower(), words[1].lower(), words[2].lower()),
-            )
+    result = await RustLayout(
+        trezorui2.select_word(
+            title="",
+            description=TR.reset__select_word_template.format(word_ordinal),
+            words=(words[0].lower(), words[1].lower(), words[2].lower()),
         )
     )
     if __debug__ and isinstance(result, str):