From b560ce53c0cd89405bb36404cd32e18efafa70c1 Mon Sep 17 00:00:00 2001
From: obrusvit <vit.obrusnik@gmail.com>
Date: Thu, 6 Jun 2024 11:58:24 +0200
Subject: [PATCH] fix(core): TS3 device freeze

[no changelog]
---
 core/src/trezor/ui/layouts/tr/__init__.py | 16 +++++++++++-----
 core/src/trezor/ui/layouts/tr/recovery.py |  4 +---
 core/src/trezor/ui/layouts/tr/reset.py    | 13 +++++--------
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/core/src/trezor/ui/layouts/tr/__init__.py b/core/src/trezor/ui/layouts/tr/__init__.py
index b27b319305..5074c9cea2 100644
--- a/core/src/trezor/ui/layouts/tr/__init__.py
+++ b/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()
diff --git a/core/src/trezor/ui/layouts/tr/recovery.py b/core/src/trezor/ui/layouts/tr/recovery.py
index 90ffe55e18..b520494f7c 100644
--- a/core/src/trezor/ui/layouts/tr/recovery.py
+++ b/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
 
 
diff --git a/core/src/trezor/ui/layouts/tr/reset.py b/core/src/trezor/ui/layouts/tr/reset.py
index 13ca822af7..2b0f4a1f73 100644
--- a/core/src/trezor/ui/layouts/tr/reset.py
+++ b/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):