diff --git a/core/src/trezor/wire/context.py b/core/src/trezor/wire/context.py
index 4dc151397c..08eaab3474 100644
--- a/core/src/trezor/wire/context.py
+++ b/core/src/trezor/wire/context.py
@@ -38,6 +38,8 @@ if TYPE_CHECKING:
 
     LoadedMessageType = TypeVar("LoadedMessageType", bound=protobuf.MessageType)
 
+    T = TypeVar("T")
+
 
 class UnexpectedMessage(Exception):
     """A message was received that is not part of the current workflow.
@@ -159,19 +161,19 @@ class Context:
 CURRENT_CONTEXT: Context | None = None
 
 
-def wait(*tasks: Awaitable) -> Any:
+def wait(task: Awaitable[T]) -> Awaitable[T]:
     """
-    Wait until one of the passed tasks finishes, and return the result, while servicing
-    the wire context.
+    Wait until the passed in task finishes, and return the result, while servicing the
+    wire context.
 
     Used to make sure the device is responsive on USB while waiting for user
-    interaction. If a message is received before any of the passed in tasks finish, it
-    raises an `UnexpectedMessage` exception, returning control to the session handler.
+    interaction. If a message is received before the task finishes, it raises an
+    `UnexpectedMessage` exception, returning control to the session handler.
     """
     if CURRENT_CONTEXT is None:
-        return loop.race(*tasks)
+        return task
     else:
-        return loop.race(CURRENT_CONTEXT.read(()), *tasks)
+        return loop.race(CURRENT_CONTEXT.read(()), task)
 
 
 async def call(