diff --git a/core/src/trezor/ui/__init__.py b/core/src/trezor/ui/__init__.py
index bf176fd641..672d1913e6 100644
--- a/core/src/trezor/ui/__init__.py
+++ b/core/src/trezor/ui/__init__.py
@@ -9,6 +9,7 @@ from trezor import io, log, loop, utils, wire, workflow
 from trezor.messages import ButtonAck, ButtonRequest
 from trezor.wire import context
 from trezorui2 import AttachType, BacklightLevels, LayoutState
+from trezor.wire.protocol_common import Context
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Generator, Generic, Iterator, TypeVar
@@ -166,7 +167,7 @@ class Layout(Generic[T]):
         self.button_request_ack_pending: bool = False
         self.transition_out: AttachType | None = None
         self.backlight_level = BacklightLevels.NORMAL
-        self.context: context.Context | None = None
+        self.context: Context | None = None
         self.state: LayoutState = LayoutState.INITIAL
 
         # Indicates whether we should use Resume attach style when launching.
@@ -217,7 +218,7 @@ class Layout(Generic[T]):
         set_current_layout(self)
 
         # save context
-        self.context = context.CURRENT_CONTEXT
+        self.context = context.get_context()
 
         # attach a timer callback and paint self
         self._event(self.layout.attach_timer_fn, self._set_timer, transition_in)