FIX: AMU pre-gate sensor names, apply_diff field preservation, code audit

  - Fix pre-gate sensor prefix: "Mmu Pre Gate N" -> "mmu_pre_gate_N"
  - Fix apply_diff losing weight_g/remaining_weight/bed_temp/mid_usage on gate rebuild
  - Add endless_spool_groups to MMUState; wired in from_status + apply_diff
  - Fix _apply_spool_data: early return on None state, flatten guard, restore bounds check
  - Fix on_load_cell_update: add startswith guard, fix float(None) on force key

Author: gmmcosta15

BlocksScreen/devices/amu/manager.py

@@ -67,6 +67,8 @@ def _apply_spool_data(self, gate: int, data: dict) -> None:
         Extracts material, color, weight from the Spoolman response dict,
         emits MMU_GATE_MAP gcode to sync Klipper, and updates the local GateInfo weight.
         """
+        if self._mmu_state is None:
+            return
         filament = data.get("filament", {})
         material = filament.get("material", "")
         color = filament.get("color_hex", "")
@@ -84,32 +86,25 @@ def _apply_spool_data(self, gate: int, data: dict) -> None:
             filament_name=filament_name,
             temperature=temperature,
         )
-        if self._mmu_state is not None:
-            if gate >= len(self._mmu_state.gates):
-                logger.warning(
-                    "Gate index %d out of range (%d gates)",
-                    gate,
-                    len(self._mmu_state.gates),
-                )
-                return
-            gates = list(self._mmu_state.gates)
-            updates = {}
-            if weight is not None:
-                updates["weight_g"] = float(weight)
-            if remaining_weight is not None:
-                updates["remaining_weight"] = float(remaining_weight)
-                self._emit_speed_gcode(gate, remaining_weight)
-            if filament_name:
-                updates["filament_name"] = filament_name
-            if temperature is not None:
-                updates["temperature"] = float(temperature)
-            if bed_temp is not None:
-                updates["bed_temp"] = int(bed_temp)
-            if updates:
-                gates[gate] = dataclasses.replace(gates[gate], **updates)
-                self._mmu_state = dataclasses.replace(
-                    self._mmu_state, gates=tuple(gates)
-                )
+        gates = list(self._mmu_state.gates)
+        if gate >= len(gates):
+            logger.warning("Gate index %d out of range (%d gate)", gate, len(gates))
+            return
+        updates = {}
+        if weight is not None:
+            updates["weight_g"] = float(weight)
+        if remaining_weight is not None:
+            updates["remaining_weight"] = float(remaining_weight)
+            self._emit_speed_gcode(gate, remaining_weight)
+        if filament_name:
+            updates["filament_name"] = filament_name
+        if temperature is not None:
+            updates["temperature"] = float(temperature)
+        if bed_temp is not None:
+            updates["bed_temp"] = int(bed_temp)
+        if updates:
+            gates[gate] = dataclasses.replace(gates[gate], **updates)
+            self._mmu_state = dataclasses.replace(self._mmu_state, gates=tuple(gates))
 
     def _emit_speed_gcode(self, gate: int, remaining_weight: float) -> None:
         """Emit MMU_GATE_MAP SPEED=x for the gate based on the spool weight profile."""
@@ -258,7 +253,7 @@ def eject_gate(self, gate: int) -> None:
         self.run_gcode_signal.emit(f"MMU_EJECT GATE={gate}")
 
     def eject_all_gates(self, num_gates: int) -> None:
-        """Fully eject filament from all gates sequentially(amu_manager.get_state().num_gates)
+        """Fully eject filament from all gates sequentially
 
         Args:
            num_gates: Total number of gates(from MMUState.num_gates)
@@ -303,10 +298,10 @@ def on_object_updated(
             self.on_load_cell_update(values, object_name)
 
     def on_pre_gate_update(self, values: dict, name: str) -> None:
-        if not name.startswith("Mmu Pre Gate "):
+        if not name.startswith("mmu_pre_gate_"):
             return
         try:
-            gate = int(name.removeprefix("Mmu Pre Gate "))
+            gate = int(name.removeprefix("mmu_pre_gate_"))
         except ValueError:
             logger.error("Failed to parse Pre-Gate: %s", name)
             return
@@ -316,15 +311,15 @@ def on_pre_gate_update(self, values: dict, name: str) -> None:
 
     def on_load_cell_update(self, values: dict, name: str) -> None:
         """Update gate weight from a Klipper load_cell sensor reading"""
-        if self._mmu_state is None:
+        if self._mmu_state is None or not name.startswith("load_cell_mmu_"):
             return
         try:
             gate = int(name.removeprefix("load_cell_mmu_"))
         except ValueError:
             logger.error("Failed parsing %s Load cell", name)
             return
 
-        weight = float(values.get("force", 0))
+        weight = float(values.get("force") or 0)
         if gate >= len(self._mmu_state.gates):
             logger.warning(
                 "Gate index %d out of range (%d gates)",

BlocksScreen/devices/amu/models.py

@@ -77,6 +77,7 @@ class MMUState:
     operation: str = ""
     sensors: dict = dataclasses.field(default_factory=dict)
     gate_speed_override: tuple[float, ...] = dataclasses.field(default_factory=tuple)
+    endless_spool_groups: tuple[int, ...] = dataclasses.field(default_factory=tuple)
 
     @property
     def is_paused(self) -> bool:
@@ -110,6 +111,10 @@ def from_status(cls, data: dict) -> "MMUState":
         filament_name = data.get("gate_filament_name", [""] * num_gates)
         temperature = data.get("gate_temperature", [None] * num_gates)
         speed_override = data.get("gate_speed_override", [])
+        weight_g = data.get("gate_weight_g", [None] * num_gates)
+        remaining_weight_list = data.get("gate_remaining_weight", [None] * num_gates)
+        bed_temperature_list = data.get("gate_bed_temp", [None] * num_gates)
+        mid_usage_list = data.get("gate_mid_usage", [False] * num_gates)
 
         gates: tuple[GateInfo, ...] = tuple(
             GateInfo(
@@ -121,6 +126,10 @@ def from_status(cls, data: dict) -> "MMUState":
                 spool_id=spool_ids[i],
                 filament_name=filament_name[i],
                 temperature=temperature[i],
+                weight_g=weight_g[i],
+                remaining_weight=remaining_weight_list[i],
+                bed_temp=bed_temperature_list[i],
+                mid_usage=mid_usage_list[i],
             )
             for i in range(num_gates)
         )
@@ -145,6 +154,7 @@ def from_status(cls, data: dict) -> "MMUState":
             operation=data.get("operation", ""),
             sensors=dict(data.get("sensors", {})),
             gate_speed_override=tuple(speed_override),
+            endless_spool_groups=tuple(data.get("endless_spool_groups", [])),
         )
 
     def gate_for_tool(self, tool: int) -> int:
@@ -187,6 +197,10 @@ def apply_diff(self, diff: dict) -> "MMUState":
                 scalar_fields["spoolman_support"] = SpoolmanSupport(
                     scalar_fields["spoolman_support"]
                 )
+            if "endless_spool_groups" in scalar_fields:
+                scalar_fields["endless_spool_groups"] = tuple(
+                    scalar_fields["endless_spool_groups"]
+                )
             return dataclasses.replace(self, **scalar_fields)
         # Gate arrays changed — need full rebuild, but we lost the raw arrays
         # Pass current gate data + diff into from_status
@@ -199,6 +213,10 @@ def apply_diff(self, diff: dict) -> "MMUState":
             "gate_filament_name": [g.filament_name for g in self.gates],
             "gate_speed_override": list(self.gate_speed_override),
             "gate_temperature": [g.temperature for g in self.gates],
+            "gate_weight_g": [g.weight_g for g in self.gates],
+            "gate_remaining_weight": [g.remaining_weight for g in self.gates],
+            "gate_bed_temp": [g.bed_temp for g in self.gates],
+            "gate_mid_usage": [g.mid_usage for g in self.gates],
         }
         merged = {**dataclasses.asdict(self), **gate_data, **diff}
         return MMUState.from_status(merged)

tests/amu/test_manager_unit.py

@@ -156,6 +156,20 @@ def test_from_status_gate_speed_override_defaults_to_empty(self, manager) -> Non
         manager.update_mmu_state(_FULL_STATUS)
         assert manager.get_state().gate_speed_override == ()
 
+    def test_from_status_parses_endless_spool_groups(self, manager) -> None:
+        data = {**_FULL_STATUS, "endless_spool_groups": [0, 1, 3,0]}
+        manager.update_mmu_state(data)
+        assert manager.get_state().endless_spool_groups == (0, 1, 3, 0)
+
+    def test_from_status_endless_spool_groups_defaults_to_empty(self, manager) -> None:
+        manager.update_mmu_state(_FULL_STATUS)
+        assert manager.get_state().endless_spool_groups == ()
+
+    def test_apply_diff_updates_endless_spool_groups(self, manager) -> None:
+        manager.update_mmu_state(_FULL_STATUS)
+        manager.update_mmu_state({"endless_spool_groups": [0, 0, 1, 1]})
+        assert manager.get_state().endless_spool_groups == (0, 0, 1, 1)
+
 
 class TestGcodeSignals:
     def test_set_gate_info(self, manager, qtbot) -> None:
@@ -265,54 +279,54 @@ class TestPregateSensors:
     def test_pre_gate_happy_path(self, manager) -> None:
         manager.update_mmu_state(_FULL_STATUS)
         data: dict[str, bool] = {"filament_detected": True}
-        manager.on_pre_gate_update(data, "Mmu Pre Gate 0")
-        manager.on_pre_gate_update(data, "Mmu Pre Gate 1")
-        manager.on_pre_gate_update(data, "Mmu Pre Gate 2")
-        manager.on_pre_gate_update(data, "Mmu Pre Gate 3")
+        manager.on_pre_gate_update(data, "mmu_pre_gate_0")
+        manager.on_pre_gate_update(data, "mmu_pre_gate_1")
+        manager.on_pre_gate_update(data, "mmu_pre_gate_2")
+        manager.on_pre_gate_update(data, "mmu_pre_gate_3")
         assert manager.get_pre_gate_sensors() == {0: True, 1: True, 2: True, 3: True}
 
     def test_pre_gate_emits_signal(self, manager, qtbot) -> None:
         with qtbot.waitSignal(manager.pre_gate_changed) as blocker:
-            manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 0")
+            manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_0")
         assert blocker.args == [0, True]
 
     def test_non_mmu_sensor_ignored(self, manager, qtbot) -> None:
         with qtbot.assertNotEmitted(manager.pre_gate_changed):
             manager.on_pre_gate_update({"filament_detected": True}, "Toolhead Sensor")
 
     def test_pre_gate_stores_state(self, manager) -> None:
-        manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 2")
+        manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_2")
         assert manager.get_pre_gate_sensors() == {2: True}
 
     def test_pre_gate_filament_not_detected(self, manager) -> None:
-        manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 0")
-        manager.on_pre_gate_update({"filament_detected": False}, "Mmu Pre Gate 1")
-        manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 2")
-        manager.on_pre_gate_update({"filament_detected": False}, "Mmu Pre Gate 3")
+        manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_0")
+        manager.on_pre_gate_update({"filament_detected": False}, "mmu_pre_gate_1")
+        manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_2")
+        manager.on_pre_gate_update({"filament_detected": False}, "mmu_pre_gate_3")
         assert manager.get_pre_gate_sensors() == {0: True, 1: False, 2: True, 3: False}
 
     def test_pre_gate_empty(self, manager) -> None:
         assert manager.get_pre_gate_sensors() == {}
 
     def test_pre_gate_multiple_emits(self, manager, qtbot) -> None:
         with qtbot.waitSignal(manager.pre_gate_changed) as blocker:
-            manager.on_pre_gate_update({"filament_detected": False}, "Mmu Pre Gate 0")
+            manager.on_pre_gate_update({"filament_detected": False}, "mmu_pre_gate_0")
         assert blocker.args == [0, False]
         with qtbot.waitSignal(manager.pre_gate_changed) as blocker:
-            manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 1")
+            manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_1")
         assert blocker.args == [1, True]
         with qtbot.waitSignal(manager.pre_gate_changed) as blocker:
-            manager.on_pre_gate_update({"filament_detected": False}, "Mmu Pre Gate 2")
+            manager.on_pre_gate_update({"filament_detected": False}, "mmu_pre_gate_2")
         assert blocker.args == [2, False]
         with qtbot.waitSignal(manager.pre_gate_changed) as blocker:
-            manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 3")
+            manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_3")
         assert blocker.args == [3, True]
 
     def test_pre_gate_update_is_dict(self, manager) -> None:
         assert isinstance(manager.get_pre_gate_sensors(), dict)
 
     def test_klippy_disconnect_clears_pre_gate_sensors(self, manager) -> None:
-        manager.on_pre_gate_update({"filament_detected": True}, "Mmu Pre Gate 0")
+        manager.on_pre_gate_update({"filament_detected": True}, "mmu_pre_gate_0")
         manager.on_klippy_state("disconnect")
         assert manager.get_pre_gate_sensors() == {}