fix variation

laser_control/controller.py

@@ -80,6 +80,11 @@ class LaserController:
         self._on_data = on_data
         self._message_id = 0
         self._last_measurements: Optional[Measurements] = None
+        # Last manual-mode params, used to restore state after stop_task()
+        self._last_temp1: float = 25.0
+        self._last_temp2: float = 25.0
+        self._last_current1: float = 30.0
+        self._last_current2: float = 30.0
 
     # ---- Connection -------------------------------------------------------
 
@@ -151,6 +156,10 @@ class LaserController:
             message_id=self._message_id,
         )
         self._send_and_read_state(cmd)
+        self._last_temp1 = validated['temp1']
+        self._last_temp2 = validated['temp2']
+        self._last_current1 = validated['current1']
+        self._last_current2 = validated['current2']
         logger.debug("Manual mode set: T1=%.2f T2=%.2f I1=%.2f I2=%.2f",
                      validated['temp1'], validated['temp2'],
                      validated['current1'], validated['current2'])
@@ -236,11 +245,34 @@ class LaserController:
                     validated['step'])
 
     def stop_task(self) -> None:
-        """Stop the current task by sending DEFAULT_ENABLE (reset)."""
-        cmd = Protocol.encode_default_enable()
-        self._send_and_read_state(cmd)
+        """Stop the current task and restore manual mode.
+
+        Sends DEFAULT_ENABLE (reset) followed by DECODE_ENABLE with the last
+        known manual-mode parameters.  This two-step sequence matches the
+        original firmware protocol: after DEFAULT_ENABLE the board is in a
+        reset state and must receive DECODE_ENABLE before it can respond to
+        TRANS_ENABLE data requests again.
+        """
+        cmd_reset = Protocol.encode_default_enable()
+        self._send_and_read_state(cmd_reset)
         logger.info("Task stopped (DEFAULT_ENABLE sent)")
 
+        # Restore manual mode so the board is ready for TRANS_ENABLE requests
+        self._message_id = (self._message_id + 1) & 0xFFFF
+        cmd_restore = Protocol.encode_decode_enable(
+            temp1=self._last_temp1,
+            temp2=self._last_temp2,
+            current1=self._last_current1,
+            current2=self._last_current2,
+            pi_coeff1_p=self._pi1_p,
+            pi_coeff1_i=self._pi1_i,
+            pi_coeff2_p=self._pi2_p,
+            pi_coeff2_i=self._pi2_i,
+            message_id=self._message_id,
+        )
+        self._send_and_read_state(cmd_restore)
+        logger.info("Manual mode restored after task stop")
+
     def get_measurements(self) -> Optional[Measurements]:
         """
         Request and return the latest measurements from the device.
@@ -339,5 +371,13 @@ class LaserController:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        # Always try to stop any running task before closing the port.
+        # If we don't, the board stays in TASK state and ignores all future
+        # commands until its power is cycled.
+        if self.is_connected:
+            try:
+                self.stop_task()
+            except Exception:
+                pass
         self.disconnect()
         return False

laser_control/example_usage.py

@@ -69,7 +69,10 @@ def example_variation_mode(port: str = None):
                 }
             )
             print("Variation task started. Collecting data for 2 s...")
-            time.sleep(2)
+            deadline = time.monotonic() + 2.0
+            while time.monotonic() < deadline:
+                ctrl.get_measurements()
+                time.sleep(0.15)
 
             ctrl.stop_task()
             print(f"Done. Collected {len(collected)} measurements.")

tests/test_integration.py

@@ -239,15 +239,19 @@ class TestConnectionManagement:
         assert "connect" in str(exc_info.value).lower()
 
     def test_stop_task_sends_default_enable(self, connected_controller, mock_serial):
-        """stop_task should send DEFAULT_ENABLE (0x2222)."""
+        """stop_task should send DEFAULT_ENABLE (0x2222) first, then DECODE_ENABLE (0x1111)."""
         mock_serial.write.reset_mock()
         connected_controller.stop_task()
 
-        assert mock_serial.write.called
-        sent_data = mock_serial.write.call_args[0][0]
-        # DEFAULT_ENABLE: 0x2222 → flipped to bytes 0x22 0x22
-        assert sent_data[0] == 0x22
-        assert sent_data[1] == 0x22
+        assert mock_serial.write.call_count >= 2
+        # First call: DEFAULT_ENABLE 0x2222 → flipped bytes 0x22 0x22
+        first_call = mock_serial.write.call_args_list[0][0][0]
+        assert first_call[0] == 0x22
+        assert first_call[1] == 0x22
+        # Second call: DECODE_ENABLE 0x1111 → flipped bytes 0x11 0x11
+        second_call = mock_serial.write.call_args_list[1][0][0]
+        assert second_call[0] == 0x11
+        assert second_call[1] == 0x11
 
     def test_reset_sends_default_enable(self, connected_controller, mock_serial):
         """reset() should also send DEFAULT_ENABLE."""