fix cmd server

BasaltTrainer/cmd_server.py

@@ -12,8 +12,9 @@ class ProtocolCode(IntEnum):
     HANDSHAKE_RESPONSE = 0x62  # Presenter -> Trainer
     DATA_READY = 0x01  # Presenter -> Trainer
     DATA_RECEIVED = 0x02  # Trainer -> Presenter
-    REQUEST_STOP = 0x71  # Presenter -> Trainer (request stop)
-    STOP = 0x71  # Trainer -> Presenter (confirm stop)
+    ACTIVELY_STOP = 0x21  # Presenter -> Trainer
+    STOP_REQUEST = 0x71  # Trainer -> Presenter
+    STOP_RESPONSE = 0x72  # Trainer -> Presenter
 
 
 class PixelKind(IntEnum):
@@ -45,42 +46,42 @@ class CmdServer:
     Command server implementation for the Trainer side according to the protocol.
     """
 
-    pipe_operator: PipeOperator
-    status: ServerStatus
+    __pipe_operator: PipeOperator
+    __status: ServerStatus
 
     def __init__(self):
-        self.pipe_operator = PipeOperator(PIPE_NAME)
-        self.status = ServerStatus.Ready
+        self.__pipe_operator = PipeOperator(PIPE_NAME)
+        self.__status = ServerStatus.Ready
 
     def __del__(self):
         """Cleanup resources when object is destroyed."""
-        self.pipe_operator.close()
+        self.__pipe_operator.close()
 
     def wait_handshake(self, payload: HandshakePayload) -> None:
         """
         Wait for handshake from Presenter, send request first and wait for response with data properties.
         Returns a tuple of (pixel_kind, width, height) from the Presenter.
         """
-        if self.status != ServerStatus.Ready:
+        if self.__status != ServerStatus.Ready:
             raise RuntimeError("unexpected server status")
 
-        # Send handshake request to Presenter (code 0x61)
-        self.pipe_operator.write(
-            CODE_PACKER.pack(ProtocolCode.HANDSHAKE_REQUEST))
+        # Send handshake request to Presenter
+        self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.HANDSHAKE_REQUEST))
         # And the payload data
-        self.pipe_operator.write(
-            HANDSHAKE_REQUEST_PACKER.pack(payload.pixel_kind, payload.width,
-                                          payload.height))
+        self.__pipe_operator.write(HANDSHAKE_REQUEST_PACKER.pack(
+            payload.pixel_kind, 
+            payload.width,
+            payload.height
+        ))
 
         # Wait for handshake response from Presenter (code 0x62)
-        code_bytes = self.pipe_operator.read(CODE_PACKER.size)
+        code_bytes = self.__pipe_operator.read(CODE_PACKER.size)
         (code, ) = CODE_PACKER.unpack(code_bytes)
         if ProtocolCode(code) != ProtocolCode.HANDSHAKE_RESPONSE:
-            raise RuntimeError(
-                "expect handshake response code, but got another")
+            raise RuntimeError("expect handshake response code, but got another")
 
         # Set status and return
-        self.status = ServerStatus.Running
+        self.__status = ServerStatus.Running
         return
 
     def tick(self, data_receiver: Callable[[], None],
@@ -89,34 +90,53 @@ class CmdServer:
         Tick function called every frame to wait for data ready from Presenter and send response.
         Returns True if a stop code was received (meaning the process should stop), False otherwise.
         """
-        if self.status != ServerStatus.Running:
+        if self.__status != ServerStatus.Running:
             raise RuntimeError("unexpected server status")
 
         # If there is stop requested, we post it first and return
         if request_stop:
-            self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP))
-            self.status = ServerStatus.Stop
+            self.__wait_stop()
             return True
 
-        # Wait for code from Presenter
-        code_bytes = self.pipe_operator.read(CODE_PACKER.size)
-        (code, ) = CODE_PACKER.unpack(code_bytes)
+        while True:
+            # Wait for code from Presenter
+            code_bytes = self.__pipe_operator.read(CODE_PACKER.size)
+            (code, ) = CODE_PACKER.unpack(code_bytes)
+
+            # Analyse code
+            match ProtocolCode(code):
+                case ProtocolCode.DATA_READY:
+                    # Receive data
+                    data_receiver()
+                    # Send data received symbol
+                    self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.DATA_RECEIVED))
+                    return False
+                case ProtocolCode.ACTIVELY_STOP:
+                    # Presenter requested stop.
+                    # Agree with it, send code and wait response
+                    self.__wait_stop()
+                    return True
+                case _:
+                    raise RuntimeError("unexpected protocol code when running")
+
+
+    def __wait_stop(self) -> None:
+        # Send stop request code
+        self.__pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP_REQUEST))
+
+        # Wait stop response code
+        while True:
+            # Accept code
+            code_bytes = self.__pipe_operator.read(CODE_PACKER.size)
+            (code, ) = CODE_PACKER.unpack(code_bytes)
+
+            # Check whether it is stop response
+            match ProtocolCode(code):
+                case ProtocolCode.STOP_RESPONSE:
+                    # Set self status and return
+                    self.__status = ServerStatus.Stop
+                    return
+                case _:
+                    raise RuntimeError("unexpected protocol code when waiting quit")
 
-        # Analyse code
-        match ProtocolCode(code):
-            case ProtocolCode.DATA_READY:
-                # Receive data
-                data_receiver()
-                # Send data received symbol
-                self.pipe_operator.write(
-                    CODE_PACKER.pack(ProtocolCode.DATA_RECEIVED))
-            case ProtocolCode.REQUEST_STOP:
-                # Presenter requested stop.
-                # Agree with it, send code and return.
-                self.pipe_operator.write(CODE_PACKER.pack(ProtocolCode.STOP))
-                self.status = ServerStatus.Stop
-                return True
-            case _:
-                raise RuntimeError("unexpected protocol code when running")
 
-        return False