feat(core): allow preempting stale THP channels on all interfaces

diff --git a/core/src/session.py b/core/src/session.py
index c8cb3215a..deeb7edea 100644
--- a/core/src/session.py
+++ b/core/src/session.py
@@ -20,14 +20,14 @@ if __debug__:
 apps.base.set_homescreen()
 workflow.start_default()
 
-# initialize the wire codec over USB
-wire.setup(usb.iface_wire)
-
 if utils.USE_BLE:
     import trezorble as ble
 
-    # initialize the wire codec over BLE
-    wire.setup(ble.interface)
+    # initialize the wire codec over USB & BLE
+    wire.setup(usb.iface_wire, ble.interface)
+else:
+    # initialize the wire codec over USB
+    wire.setup(usb.iface_wire)
 
 # start the event loop
 loop.run()
diff --git a/core/src/trezor/wire/__init__.py b/core/src/trezor/wire/__init__.py
index a2565d25b..3c42f4439 100644
--- a/core/src/trezor/wire/__init__.py
+++ b/core/src/trezor/wire/__init__.py
@@ -76,11 +76,6 @@ class Provider(Generic[T]):
         return obj
 
 
-def setup(iface: WireInterface) -> None:
-    """Initialize the wire stack on the provided WireInterface."""
-    loop.schedule(handle_session(iface))
-
-
 if utils.USE_THP:
     from .thp.memory_manager import ThpBuffer
 
@@ -99,21 +94,25 @@ if utils.USE_THP:
                     return result
             return None
 
-    async def handle_session(iface: WireInterface) -> None:
-        ctx = ThpContext(iface)
+    def setup(*ifaces: WireInterface) -> None:
+        """Initialize the wire stack on the provided interfaces."""
+        loop.schedule(handle_session_thp(*ifaces))
+
+    async def handle_session_thp(*ifaces: WireInterface) -> None:
+        ctx = ThpContext(*ifaces)
         if __debug__:
-            _THP_CHANNELS.append(ctx._channels)
+            _THP_CHANNELS.extend(iface_ctx._channels for iface_ctx in ctx._iface_ctxs)
+
         try:
             while (channel := await ctx.get_next_message()) is None:
-                if __debug__:
-                    # happens if another interface is active and using THP buffers.
-                    log.error(__name__, "Another interface is active", iface=iface)
+                # wait until a new channel is established (on any interface)
+                pass
 
             while await received_message_handler.handle_received_message(channel):
                 pass
         finally:
             if __debug__:
-                log.debug(__name__, "Finished THP session", iface=iface)
+                log.debug(__name__, "Finished THP session: %s", ifaces)
             # Wait for all active workflows to finish.
             await workflow.join_all()
             if __debug__:
@@ -128,7 +127,12 @@ else:
     # Acquired by the first call to `CodecContext.read_from_wire()`.
     WIRE_BUFFER_PROVIDER = Provider(bytearray(8192))
 
-    async def handle_session(iface: WireInterface) -> None:
+    def setup(*ifaces: WireInterface) -> None:
+        """Initialize the wire stack on the provided interfaces."""
+        for iface in ifaces:
+            loop.schedule(handle_session_codec(iface))
+
+    async def handle_session_codec(iface: WireInterface) -> None:
         ctx = CodecContext(iface, WIRE_BUFFER_PROVIDER)
         next_msg: protocol_common.Message | None = None
 
diff --git a/core/src/trezor/wire/thp/channel.py b/core/src/trezor/wire/thp/channel.py
index 9d90fcf3c..690ca253d 100644
--- a/core/src/trezor/wire/thp/channel.py
+++ b/core/src/trezor/wire/thp/channel.py
@@ -51,7 +51,7 @@ if TYPE_CHECKING:
     from trezor.messages import ThpPairingCredential
     from trezor.wire import WireInterface
 
-    from .interface_context import ThpContext
+    from .interface_context import InterfaceContext
     from .memory_manager import ThpBuffer
     from .pairing_context import PairingContext
     from .session_context import GenericSessionContext
@@ -155,14 +155,14 @@ class Channel:
     def __init__(
         self,
         channel_cache: ChannelCache,
-        ctx: ThpContext,
+        ctx: InterfaceContext,
         buffers: tuple[ThpBuffer, ThpBuffer],
     ) -> None:
         assert ctx._iface.iface_num() == channel_cache.get_int(CHANNEL_IFACE)
 
         # Channel properties
         self.channel_id: bytes = channel_cache.channel_id
-        self.ctx: ThpContext = ctx
+        self.iface_ctx: InterfaceContext = ctx
         self.read_buf, self.write_buf = buffers
         if __debug__:
             self._log("channel initialization")
@@ -179,7 +179,7 @@ class Channel:
 
     @property
     def iface(self) -> WireInterface:
-        return self.ctx._iface
+        return self.iface_ctx._iface
 
     def clear(self) -> None:
         clear_sessions_with_channel_id(self.channel_id)
@@ -286,11 +286,24 @@ class Channel:
         self, timeout_ms: int | None = None
     ) -> memoryview:
         """Doesn't block if a message has been already reassembled."""
+        thp_ctx = self.iface_ctx.thp_ctx
         while self.reassembler.message is None:
-            # receive and reassemble a new message from this channel
-            channel = await self.ctx.get_next_message(timeout_ms=timeout_ms)
-            if channel is None:
-                self.preempt_if_stale()
+            # receive and reassemble a new message from any THP channel
+            try:
+                channel = await thp_ctx.get_next_message(timeout_ms=timeout_ms)
+                if channel is None:
+                    continue
+            except ChannelPreemptedException:
+                elapsed_ms = utime.ticks_diff(utime.ticks_ms(), self.last_write_ms)
+                # allow preempting channel only after enough time has passed
+                is_stale = elapsed_ms > _PREEMPT_TIMEOUT_MS
+                if __debug__:
+                    self._log(
+                        f"Interrupted channel after {elapsed_ms} ms",
+                        logger=(log.error if is_stale else log.warning),
+                    )
+                if is_stale:
+                    raise
                 continue
 
             if channel is self:
@@ -396,15 +409,6 @@ class Channel:
     ) -> Awaitable[None]:
         return self.write_encrypted_payload(ctrl_byte, payload)
 
-    def preempt_if_stale(self) -> None:
-        elapsed_ms = utime.ticks_diff(utime.ticks_ms(), self.last_write_ms)
-        preempt = elapsed_ms > _PREEMPT_TIMEOUT_MS
-        if __debug__:
-            logger = log.error if preempt else log.warning
-            self._log(f"Interrupted channel after {elapsed_ms} ms", logger=logger)
-        if preempt:
-            raise ChannelPreemptedException
-
     async def write_encrypted_payload(self, ctrl_byte: int, payload: bytes) -> None:
         if __debug__:
             self._log("write_encrypted_payload_loop")
@@ -423,7 +427,9 @@ class Channel:
         self.last_write_ms = utime.ticks_ms()
 
         for i in range(_MAX_RETRANSMISSION_COUNT):
-            result = await race(self.ctx.write_payload(header, payload), _WRITE_TIMEOUT)
+            result = await race(
+                self.iface_ctx.write_payload(header, payload), _WRITE_TIMEOUT
+            )
             if isinstance(result, int):
                 if __debug__:
                     log.error(__name__, "Sending is stuck for %d ms", _WRITE_TIMEOUT_MS)
@@ -492,7 +498,7 @@ def send_ack(channel: Channel, ack_bit: int) -> Awaitable[None]:
             ack_bit,
             iface=channel.iface,
         )
-    return channel.ctx.write_payload(header, b"")
+    return channel.iface_ctx.write_payload(header, b"")
 
 
 def handle_ack(ctx: Channel, ack_bit: int) -> None:
diff --git a/core/src/trezor/wire/thp/interface_context.py b/core/src/trezor/wire/thp/interface_context.py
index 22bbc7f26..328f27dfc 100644
--- a/core/src/trezor/wire/thp/interface_context.py
+++ b/core/src/trezor/wire/thp/interface_context.py
@@ -8,7 +8,8 @@ from storage.cache_thp import (
     find_allocated_channel,
     update_channel_last_used,
 )
-from trezor import io, loop, utils
+from trezor import io, utils
+from trezor.loop import Timeout, race, sleep, wait
 
 from . import (
     CHANNEL_ALLOCATION_REQ,
@@ -22,7 +23,7 @@ from . import (
     control_byte,
     get_channel_allocation_response,
 )
-from .channel import Channel
+from .channel import Channel, ChannelPreemptedException
 from .checksum import CHECKSUM_LENGTH
 
 if __debug__:
@@ -30,65 +31,100 @@ if __debug__:
 
 if TYPE_CHECKING:
     from trezorio import WireInterface
-    from typing import Awaitable, Iterable
+    from typing import Awaitable, Generator, Iterable, NoReturn
 
 _BROADCAST_PAYLOAD_LENGTH = const(12)
 
 
+# Uses `yield` instead of `await` to avoid allocations.
+def _timeout_after(ms: int) -> Generator[sleep, int, NoReturn]:
+    yield sleep(ms)
+    raise Timeout
+
+
 class ThpContext:
+    def __init__(self, *ifaces: WireInterface) -> None:
+        max_packet_len = max(iface.RX_PACKET_LEN for iface in ifaces)
+        self._packet_buf = bytearray(max_packet_len)
+        self._packet_view = memoryview(self._packet_buf)
+        self._iface_ctxs = [InterfaceContext(iface, self) for iface in ifaces]
+
+    async def get_next_message(self, timeout_ms: int | None = None) -> Channel | None:
+        """
+        Reassemble a valid THP payload from any THP interface, and return its channel.
+
+        Also handle THP channel allocation.
+        """
+        # wait until one of the channels becomes readable
+        children = (iface_ctx._wait_for_packet() for iface_ctx in self._iface_ctxs)
+        if timeout_ms is None:
+            race_task = race(*children)
+        else:
+            race_task = race(*children, _timeout_after(timeout_ms))
+
+        (iface_ctx, packet_len) = await race_task  # will raise on timeout
+        assert packet_len == iface_ctx._iface.RX_PACKET_LEN
+
+        # read and handle the packet using its `InterfaceContext`
+        iface_ctx._iface.read(self._packet_buf, 0)
+        return await iface_ctx.handle_packet(self._packet_view[:packet_len])
+
+
+class InterfaceContext:
     """
-    This class allows fetching multi-packet THP payloads from a given interface.
+    This class handles multi-packet THP payloads from a single interface.
     It also handles and responds to low-level single packet THP messages, creating new channels if needed.
     """
 
-    def __init__(self, iface: WireInterface) -> None:
+    def __init__(self, iface: WireInterface, thp_ctx: ThpContext) -> None:
         self._iface = iface
-        self._read = loop.wait(iface.iface_num() | io.POLL_READ)
-        self._write = loop.wait(iface.iface_num() | io.POLL_WRITE)
+        self._read = wait(iface.iface_num() | io.POLL_READ)
+        self._write = wait(iface.iface_num() | io.POLL_WRITE)
         self._channels: dict[int, Channel] = {}
+        self.thp_ctx = thp_ctx
 
-    async def get_next_message(self, timeout_ms: int | None = None) -> Channel | None:
+    def _wait_for_packet(self) -> Generator[wait, int, tuple["InterfaceContext", int]]:
+        """Block until this interface is readable.
+
+        It adapts `loop.wait`, to be used in a `race()` over multiple THP interfaces by `ThpContext.get_next_message()`.
         """
-        Reassemble a valid THP payload and return its channel.
+        # Uses `yield` instead of `await` to avoid allocations.
+        packet_len = yield self._read
+        return self, packet_len
+
+    async def handle_packet(self, packet: memoryview) -> Channel | None:
+        """
+        Reassemble a valid THP payload and return its channel, if reassembly succeeds.
+        Otherwise, returns `None` and should be called again (with the next packet).
 
         Also handle THP channel allocation.
         """
-        from .. import THP_BUFFERS_PROVIDER
-
-        packet = bytearray(self._iface.RX_PACKET_LEN)
-        while True:
-            self._read.timeout_ms = timeout_ms
-            packet_len = await self._read
-            assert packet_len is not None
-            assert packet_len == len(packet)
-            self._iface.read(packet, 0)
-
-            ctrl_byte = _get_ctrl_byte(packet)
-            if ctrl_byte == CODEC_V1:
-                await self._handle_codec_v1(packet)
-                continue
-
-            cid = ustruct.unpack(">BH", packet)[1]
-            if cid == BROADCAST_CHANNEL_ID:
-                await self._handle_broadcast(packet)
-                continue
-
-            if (cache := find_allocated_channel(cid)) is None:
-                if not control_byte.is_continuation(_get_ctrl_byte(packet)):
-                    await self.write_error(cid, ThpErrorType.UNALLOCATED_CHANNEL)
-                continue
-
-            if (channel := self._channels.get(cid)) is None:
-                if (buffers := THP_BUFFERS_PROVIDER.take()) is None:
-                    # concurrent payload reassembly is not supported
-                    await self.write_error(cid, ThpErrorType.TRANSPORT_BUSY)
-                    return None  # try to preempt the caller (if stale)
-
-                channel = self._channels[cid] = Channel(cache, self, buffers)
-
-            if channel.reassemble(packet):
-                update_channel_last_used(channel.channel_id)
-                return channel
+        ctrl_byte = _get_ctrl_byte(packet)
+        if ctrl_byte == CODEC_V1:
+            return await self._handle_codec_v1(packet)
+
+        cid = ustruct.unpack(">BH", packet)[1]
+        if cid == BROADCAST_CHANNEL_ID:
+            return await self._handle_broadcast(packet)
+
+        if (cache := find_allocated_channel(cid)) is None:
+            if not control_byte.is_continuation(_get_ctrl_byte(packet)):
+                await self.write_error(cid, ThpErrorType.UNALLOCATED_CHANNEL)
+            return None
+
+        if (channel := self._channels.get(cid)) is None:
+            from .. import THP_BUFFERS_PROVIDER
+
+            if (buffers := THP_BUFFERS_PROVIDER.take()) is None:
+                # concurrent payload reassembly is not supported
+                await self.write_error(cid, ThpErrorType.TRANSPORT_BUSY)
+                raise ChannelPreemptedException  # try to preempt the caller (if stale)
+
+            channel = self._channels[cid] = Channel(cache, self, buffers)
+
+        if channel.reassemble(packet):
+            update_channel_last_used(channel.channel_id)
+            return channel
 
     def write_payload(self, header: PacketHeader, payload: bytes) -> Awaitable[None]:
         checksum = crc.crc32(payload, crc.crc32(header.to_bytes()))
diff --git a/core/src/trezor/wire/thp/received_message_handler.py b/core/src/trezor/wire/thp/received_message_handler.py
index bef6b6598..0240c68cf 100644
--- a/core/src/trezor/wire/thp/received_message_handler.py
+++ b/core/src/trezor/wire/thp/received_message_handler.py
@@ -69,12 +69,12 @@ async def handle_received_message(channel: Channel) -> bool:
         error_message = Failure(code=FailureType.ThpUnallocatedSession)
         await channel.write(error_message, e.session_id)
     except ThpDecryptionError:
-        await channel.ctx.write_error(
+        await channel.iface_ctx.write_error(
             channel.get_channel_id_int(), ThpErrorType.DECRYPTION_FAILED
         )
         channel.clear()
     except ThpDeviceLockedError:
-        await channel.ctx.write_error(
+        await channel.iface_ctx.write_error(
             channel.get_channel_id_int(), ThpErrorType.DEVICE_LOCKED
         )
     return False
diff --git a/core/tests/test_trezor.wire.thp.py b/core/tests/test_trezor.wire.thp.py
deleted file mode 100644
index de6bff166..000000000
--- a/core/tests/test_trezor.wire.thp.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# flake8: noqa: F403,F405
-from common import *  # isort:skip
-from mock_wire_interface import MockHID
-from trezor import io
-
-if utils.USE_THP:
-    import thp_common
-    from trezor.wire import handle_session as thp_main_loop
-
-
-@unittest.skipUnless(utils.USE_THP, "only needed for THP")
-class TestTrezorHostProtocol(unittest.TestCase):
-
-    def __init__(self):
-        if __debug__:
-            thp_common.suppress_debug_log()
-        super().__init__()
-
-    def setUp(self):
-        self.interface = MockHID()
-
-    def test_codec_message(self):
-        self.assertEqual(len(self.interface.data), 0)
-        gen = thp_main_loop(self.interface)
-        gen.send(None)
-
-        # There should be a failiure response to received init packet (starts with "?##")
-        test_codec_message = b"?## Some data"
-        self.interface.mock_read(test_codec_message, gen)
-        gen.send(None)
-        self.assertEqual(len(self.interface.data), 1)
-
-        expected_response = b"?##\x00\x03\x00\x00\x00\x14\x08\x11"
-        self.assertEqual(
-            self.interface.data[-1][: len(expected_response)], expected_response
-        )
-
-        # There should be no response for continuation packet (starts with "?" only)
-        test_codec_message_2 = b"? Cont packet"
-        self.interface.mock_read(test_codec_message_2, gen)
-
-        # Check that sending None fails on AssertionError
-        with self.assertRaises(AssertionError):
-            gen.send(None)
-        self.assertEqual(len(self.interface.data), 1)
-
-    def test_message_on_unallocated_channel(self):
-        gen = thp_main_loop(self.interface)
-        query = gen.send(None)
-        self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
-        message_to_channel_789a = (
-            b"\x04\x78\x9a\x00\x0c\x00\x11\x22\x33\x44\x55\x66\x77\x96\x64\x3c\x6c"
-        )
-        self.interface.mock_read(message_to_channel_789a, gen)
-        gen.send(None)
-        unallocated_chanel_error_on_channel_789a = "42789a0005027b743563000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
-        self.assertEqual(
-            utils.hexlify_if_bytes(self.interface.data[-1]),
-            unallocated_chanel_error_on_channel_789a,
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/core/tests/test_trezor.wire.thp.writer.py b/core/tests/test_trezor.wire.thp.writer.py
index bbbdea4e1..b26159bc0 100644
--- a/core/tests/test_trezor.wire.thp.writer.py
+++ b/core/tests/test_trezor.wire.thp.writer.py
@@ -77,7 +77,8 @@ class TestTrezorHostProtocolWriter(unittest.TestCase):
 
     def setUp(self):
         self.interface = MockHID()
-        self.ctx = ThpContext(self.interface)
+        thp_ctx = ThpContext(self.interface)
+        (self.ctx,) = thp_ctx._iface_ctxs
 
     def test_write_empty_payload(self):
         header = PacketHeader(ENCRYPTED, 4660, 4)
diff --git a/core/tests/thp_common.py b/core/tests/thp_common.py
index 85b12da35..b40371287 100644
--- a/core/tests/thp_common.py
+++ b/core/tests/thp_common.py
@@ -18,18 +18,17 @@ if utils.USE_THP:
 
     def prepare_context() -> None:
         mock_iface = MockHID()
-        channel_cache = create_new_channel(mock_iface)
+        channel = get_new_channel(mock_iface)
         session_cache = cache_thp.create_or_replace_session(
-            channel_cache, session_id=b"\x01"
-        )
-        channel = Channel(
-            channel_cache, ThpContext(mock_iface), (ThpBuffer(), ThpBuffer())
+            channel.channel_cache, session_id=b"\x01"
         )
         context.CURRENT_CONTEXT = SessionContext(channel, session_cache)
 
     def get_new_channel(iface: WireInterface) -> Channel:
         channel_cache = create_new_channel(iface)
-        return Channel(channel_cache, ThpContext(iface), (ThpBuffer(), ThpBuffer()))
+        thp_ctx = ThpContext(iface)
+        (iface_ctx,) = thp_ctx._iface_ctxs
+        return Channel(channel_cache, iface_ctx, (ThpBuffer(), ThpBuffer()))
 
 
 if __debug__: