From 5bd51bc6b72ec88e8d329a96872ffce3c523c6fb Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Mon, 6 Jan 2025 16:06:50 +0100
Subject: [PATCH] TEMP: try fix mailbox issue

---
 core/src/apps/base.py                         |  3 ++-
 .../wire/thp/received_message_handler.py      | 27 ++++++++++---------
 core/src/trezor/wire/thp/session_context.py   |  4 ++-
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/core/src/apps/base.py b/core/src/apps/base.py
index 86ffd8af01..bebc5822a0 100644
--- a/core/src/apps/base.py
+++ b/core/src/apps/base.py
@@ -269,7 +269,8 @@ if utils.USE_THP:
             )
 
         channel.sessions[new_session.session_id] = new_session
-        loop.schedule(new_session.handle())
+        new_session.handle_spawn = loop.spawn(new_session.handle())
+        loop.schedule(new_session.handle_spawn)
 
         return Success(message="New session created.")
 
diff --git a/core/src/trezor/wire/thp/received_message_handler.py b/core/src/trezor/wire/thp/received_message_handler.py
index c2d39dbd1e..a992ee856e 100644
--- a/core/src/trezor/wire/thp/received_message_handler.py
+++ b/core/src/trezor/wire/thp/received_message_handler.py
@@ -382,7 +382,8 @@ async def _handle_state_ENCRYPTED_TRANSPORT(ctx: Channel, message_length: int) -
             s = SeedlessSessionContext(ctx, session_id)
 
         ctx.sessions[session_id] = s
-        loop.schedule(s.handle())
+        s.handle_spawn = loop.spawn(s.handle())
+        loop.schedule(s.handle_spawn)
 
     elif ctx.sessions[session_id].get_session_state() is SessionState.UNALLOCATED:
         raise ThpUnallocatedSessionError(session_id)
@@ -390,18 +391,20 @@ async def _handle_state_ENCRYPTED_TRANSPORT(ctx: Channel, message_length: int) -
     s = ctx.sessions[session_id]
     update_session_last_used(s.channel_id, (s.session_id).to_bytes(1, "big"))
 
-    s.incoming_message.put(
-        Message(
-            message_type,
-            buffer[
-                INIT_HEADER_LENGTH
-                + MESSAGE_TYPE_LENGTH
-                + SESSION_ID_LENGTH : message_length
-                - CHECKSUM_LENGTH
-                - TAG_LENGTH
-            ],
-        )
+    message = Message(
+        message_type,
+        buffer[
+            INIT_HEADER_LENGTH
+            + MESSAGE_TYPE_LENGTH
+            + SESSION_ID_LENGTH : message_length
+            - CHECKSUM_LENGTH
+            - TAG_LENGTH
+        ],
     )
+    if not s.incoming_message.is_empty():
+        await s.handle_spawn
+        s.incoming_message.clear()
+    s.incoming_message.put(message)
 
 
 async def _handle_pairing(ctx: Channel, message_length: int) -> None:
diff --git a/core/src/trezor/wire/thp/session_context.py b/core/src/trezor/wire/thp/session_context.py
index 8c76136ec4..ff8b2eb6d8 100644
--- a/core/src/trezor/wire/thp/session_context.py
+++ b/core/src/trezor/wire/thp/session_context.py
@@ -34,6 +34,7 @@ class GenericSessionContext(Context):
         self.channel: Channel = channel
         self.session_id: int = session_id
         self.incoming_message = loop.mailbox()
+        self.handle_spawn: loop.spawn | None = None
 
     async def handle(self) -> None:
         if __debug__ and utils.ALLOW_DEBUG_MESSAGES:
@@ -51,7 +52,8 @@ class GenericSessionContext(Context):
             next_message = None
             try:
                 if await self._handle_message(message):
-                    loop.schedule(self.handle())
+                    self.handle_spawn = loop.spawn(self.handle())
+                    loop.schedule(self.handle_spawn)
                     return
             except UnexpectedMessageException as unexpected:
                 # The workflow was interrupted by an unexpected message. We need to