fix(usage-tracker): split mid-print AMS-Backup spool switch correctly (#1771)

Reporter forcefully started a print needing ~260 g with 180 g on the
  first spool and a backup spool in the AMS. Printer correctly consumed
  spool 1, AMS Backup switched, spool 2 finished the print. Bambuddy
  attributed all 260 g to spool 2 -- spool 1 untouched in inventory.

  Two stacking bugs produced the exact "all to second spool" symptom for
  prints without per-layer 3MF gcode data:

  1. bambu_mqtt.py:2135 wrote state.total_layers = int(data["total_layer_num"])
     unconditionally. P1S firmware pushes total_layer_num=0 at print end
     (same reset pattern other models do for layer_num / progress). The
     unconditional write clobbered the slicer's actual total to 0 before
     the usage tracker read it.

  2. usage_tracker.py:1129-1137 linear-fallback dumped EVERYTHING onto the
     last segment when total_layers was 0:
       if total_layers > 0:
           segment_grams = total_weight * (seg_end_layer - seg_start_layer) / total_layers
       else:
           segment_grams = 0.0   # <- entire print weight ends up on last segment

     Path 2 (AMS remain% delta) couldn't recover because (a) the emptied
     spool reported remain=-1 and (b) Bug-A had already added the second
     spool's key to handled_trays, suppressing the Path 2 lookup.

  Fix:

  - bambu_mqtt.py: only overwrite state.total_layers when the incoming
    value is positive (mirror of the existing _last_valid_layer_num
    pattern at line 2127). Explicit reset on new print start at
    _handle_print_start so the previous print's total can't bleed in.

  - usage_tracker.py: cascade the linear-fallback denominator -
    state.total_layers, then last_layer_num (already threaded in for
    the last_progress fallback), then equal-split as a bounded fence.
    Equal-split is still wrong but never dumps the whole print on the
    last segment, which was strictly worse.

backend/app/services/bambu_mqtt.py

@@ -2133,7 +2133,14 @@ class BambuMQTTClient:
             if new_layer > old_layer and self.on_layer_change:
                 self.on_layer_change(new_layer)
         if "total_layer_num" in data:
-            self.state.total_layers = int(data["total_layer_num"])
+            # Some firmware (P1S observed) resets `total_layer_num` to 0 at
+            # print end — same shape as the `layer_num` reset guarded above.
+            # Preserve the last known good value so the usage-tracker split
+            # path (#1771) has a denominator that survives the reset frame.
+            # Explicit reset to 0 happens on print start (`_handle_print_start`).
+            new_total = int(data["total_layer_num"])
+            if new_total > 0:
+                self.state.total_layers = new_total
 
         # Fan speeds (MQTT sends as string "0"-"15" representing speed levels, or percentage)
         # Convert to 0-100 percentage for display
@@ -3120,6 +3127,12 @@ class BambuMQTTClient:
             self.state.hms_errors = []
             # Reset layer tracking for new print (needed for layer-based timelapse)
             self.state.layer_num = 0
+            # Reset total_layers so the previous print's value can't bleed into
+            # this print's usage-tracker split before the new push_status arrives
+            # with the slicer's total (#1771 follow-on to the preservation guard
+            # above at line ~2135 — the guard now ignores firmware-reset 0s, so
+            # the explicit reset has to happen here instead).
+            self.state.total_layers = 0
             # Reset completion tracking for new print
             self._was_running = True
             self._completion_triggered = False

backend/app/services/usage_tracker.py

@@ -1127,14 +1127,27 @@ async def _track_from_3mf(
                     mm_at_end = get_cumulative_usage_at_layer(split_layer_usage, seg_end_layer).get(filament_id, 0)
                     segment_grams = mm_to_grams(mm_at_end - mm_at_start, diameter, density)
                 else:
-                    # No per-layer data: linear fallback by layer ratio
+                    # No per-layer data: linear fallback by layer ratio (#1771).
+                    # Cascade denominators because firmware on some models (P1S
+                    # observed) resets `total_layer_num` to 0 at print end —
+                    # `last_layer_num` is the print's last-valid layer captured
+                    # mid-print and survives that reset (same shape as the
+                    # `last_progress` fallback at line 1040). Equal-split is the
+                    # last-resort fence: still wrong, but bounded — never dumps
+                    # the entire print onto the last segment, which was the
+                    # original #1771 symptom for the reporter (P1S, AMS Backup
+                    # fed from spool 1 then spool 2, all 260 g credited to
+                    # spool 2 even though spool 1 had given up its 180 g).
                     seg_end_layer = tray_changes[seg_idx + 1][1]
-                    total_layers = state.total_layers if state else 0
-                    if total_layers > 0:
-                        segment_grams = total_weight * (seg_end_layer - seg_start_layer) / total_layers
+                    denom = (state.total_layers if state else 0) or last_layer_num
+                    if denom > 0:
+                        segment_grams = total_weight * (seg_end_layer - seg_start_layer) / denom
                     else:
-                        # Can't compute ratio — assign all to last segment
-                        segment_grams = 0.0
+                        # No layer information available from any source —
+                        # spread evenly across segments. The last segment will
+                        # get the rounding remainder via the `is_last` branch
+                        # above on its own iteration.
+                        segment_grams = total_weight / len(tray_changes)
 
                 sum_previous += segment_grams
                 if segment_grams <= 0:

backend/tests/unit/services/test_bambu_mqtt.py

@@ -5795,6 +5795,65 @@ class TestPrintRunningObservedCallback:
         }
 
 
+class TestTotalLayersPreservation:
+    """#1771: P1S firmware resets `total_layer_num` to 0 at print end. Without
+    this guard, the usage tracker's split path saw `state.total_layers = 0` at
+    completion and dumped the whole print onto the last spool.
+
+    These tests pin the preservation pattern (mirror of `_last_valid_layer_num`)
+    and the explicit reset on new print start so the previous print's total
+    can't bleed into the next.
+    """
+
+    @pytest.fixture
+    def mqtt_client(self):
+        from backend.app.services.bambu_mqtt import BambuMQTTClient
+
+        client = BambuMQTTClient(
+            ip_address="192.168.1.100",
+            serial_number="TEST123",
+            access_code="12345678",
+        )
+        return client
+
+    def test_nonzero_total_layer_num_sets_state(self, mqtt_client):
+        # Baseline: a fresh push with the slicer's total updates state.total_layers.
+        mqtt_client._process_message({"print": {"total_layer_num": 260}})
+        assert mqtt_client.state.total_layers == 260
+
+    def test_zero_total_layer_num_does_not_clobber_cached_value(self, mqtt_client):
+        # Firmware-reset frame: total_layer_num=0 arrives mid- or end-of-print.
+        # The guard must NOT overwrite the previously-captured 260.
+        mqtt_client._process_message({"print": {"total_layer_num": 260}})
+        mqtt_client._process_message({"print": {"total_layer_num": 0}})
+        assert mqtt_client.state.total_layers == 260
+
+    def test_print_start_explicitly_resets_total_layers(self, mqtt_client):
+        # Without the explicit reset on print start, the previous print's total
+        # would persist into the new print until its first total_layer_num push
+        # arrived — which is exactly the kind of cross-print bleed the
+        # preservation guard above otherwise opens up.
+        mqtt_client._process_message({"print": {"total_layer_num": 260}})
+        assert mqtt_client.state.total_layers == 260
+
+        # Simulate the new-print-start trigger shape (is_new_print path):
+        # state was previously RUNNING on an old file; now we observe a
+        # different file going RUNNING.
+        mqtt_client._previous_gcode_state = "RUNNING"
+        mqtt_client._previous_gcode_file = "/data/Metadata/old_print.gcode"
+        mqtt_client._was_running = True
+        mqtt_client._process_message(
+            {
+                "print": {
+                    "gcode_state": "RUNNING",
+                    "gcode_file": "/data/Metadata/new_print.gcode",
+                    "subtask_name": "new_print",
+                }
+            }
+        )
+        assert mqtt_client.state.total_layers == 0
+
+
 class TestAmsFilamentBackupHoldTimer:
     """Regression: stale push_status arriving within the hold window after a
     toggle command MUST NOT flip ams_filament_backup back to the printer's

backend/tests/unit/test_usage_tracker.py

@@ -1415,6 +1415,141 @@ class TestTrayChangeSplit:
         assert results[2]["ams_id"] == 0
         assert results[2]["tray_id"] == 2
 
+    @pytest.mark.asyncio
+    async def test_tray_switch_uses_last_layer_num_when_total_layers_reset(self):
+        """#1771 regression: P1S firmware resets `total_layer_num` to 0 at print
+        end; without the cascade the linear fallback collapsed to `0.0` per
+        non-last segment and dumped the whole print onto the last spool. With
+        the fix, `last_layer_num` (the print's last-valid layer captured before
+        the firmware reset) is the substitute denominator.
+
+        Reporter's exact shape: print needed ~260 g, started on a 180 g spool,
+        AMS Backup switched at ~70% through, second spool finished the print.
+        Before fix: spool 1 → 0 g, spool 2 → 260 g (the bug).
+        After fix:  spool 1 → 180 g, spool 2 → 80 g (correct).
+        """
+        spool_a = _make_spool(spool_id=10, label_weight=1000)
+        spool_b = _make_spool(spool_id=20, label_weight=1000)
+        assign_a = _make_assignment(spool_id=10, ams_id=0, tray_id=0)
+        assign_b = _make_assignment(spool_id=20, ams_id=0, tray_id=1)
+        archive = _make_archive(archive_id=171)
+
+        db = _mock_db_sequential([archive, None, assign_a, spool_a, assign_b, spool_b])
+
+        # Firmware reset: state.total_layers is 0 by the time usage_tracker runs.
+        # last_layer_num threaded in from on_print_complete is the survival value.
+        printer_manager = MagicMock()
+        printer_manager.get_status.return_value = SimpleNamespace(
+            progress=100,
+            layer_num=0,  # also reset
+            tray_now=1,
+            last_loaded_tray=1,
+            total_layers=0,  # the bug trigger
+            tray_change_log=[(0, 0), (1, 180)],  # switched at layer 180 of 260
+        )
+
+        filament_usage = [{"slot_id": 1, "used_g": 260.0, "type": "PLA", "color": ""}]
+        handled_trays: set[tuple[int, int]] = set()
+
+        with (
+            patch("backend.app.core.config.settings") as mock_settings,
+            patch(
+                "backend.app.utils.threemf_tools.extract_filament_usage_from_3mf",
+                return_value=filament_usage,
+            ),
+            patch(
+                "backend.app.utils.threemf_tools.extract_layer_filament_usage_from_3mf",
+                return_value=None,  # No per-layer 3MF data — force linear fallback path
+            ),
+        ):
+            mock_settings.base_dir = MagicMock()
+            mock_path = MagicMock()
+            mock_path.exists.return_value = True
+            mock_settings.base_dir.__truediv__ = MagicMock(return_value=mock_path)
+
+            results = await _track_from_3mf(
+                printer_id=1,
+                archive_id=171,
+                status="completed",
+                print_name="#1771 repro",
+                handled_trays=handled_trays,
+                printer_manager=printer_manager,
+                db=db,
+                last_layer_num=260,  # survives the firmware reset of total_layer_num
+            )
+
+        # Both segments must be attributed correctly.
+        assert len(results) == 2
+        # Segment 1: tray 0, layers 0-180 of 260 → 260 * 180/260 = 180.0 g
+        assert results[0]["ams_id"] == 0
+        assert results[0]["tray_id"] == 0
+        assert results[0]["weight_used"] == 180.0
+        # Segment 2: tray 1, remainder = 260 - 180 = 80.0 g
+        assert results[1]["ams_id"] == 0
+        assert results[1]["tray_id"] == 1
+        assert results[1]["weight_used"] == 80.0
+
+    @pytest.mark.asyncio
+    async def test_tray_switch_equal_split_when_no_layer_info_at_all(self):
+        """Defensive fence: when neither `state.total_layers` nor `last_layer_num`
+        survives (older firmware / edge case), equal-split across segments is the
+        last-resort fallback. Still wrong but BOUNDED — the original bug dumped
+        the whole print weight onto the last segment, which was strictly worse.
+        """
+        spool_a = _make_spool(spool_id=10, label_weight=1000)
+        spool_b = _make_spool(spool_id=20, label_weight=1000)
+        assign_a = _make_assignment(spool_id=10, ams_id=0, tray_id=0)
+        assign_b = _make_assignment(spool_id=20, ams_id=0, tray_id=1)
+        archive = _make_archive(archive_id=172)
+
+        db = _mock_db_sequential([archive, None, assign_a, spool_a, assign_b, spool_b])
+
+        printer_manager = MagicMock()
+        printer_manager.get_status.return_value = SimpleNamespace(
+            progress=100,
+            layer_num=0,
+            tray_now=1,
+            last_loaded_tray=1,
+            total_layers=0,  # neither source available
+            tray_change_log=[(0, 0), (1, 50)],
+        )
+
+        filament_usage = [{"slot_id": 1, "used_g": 60.0, "type": "PLA", "color": ""}]
+        handled_trays: set[tuple[int, int]] = set()
+
+        with (
+            patch("backend.app.core.config.settings") as mock_settings,
+            patch(
+                "backend.app.utils.threemf_tools.extract_filament_usage_from_3mf",
+                return_value=filament_usage,
+            ),
+            patch(
+                "backend.app.utils.threemf_tools.extract_layer_filament_usage_from_3mf",
+                return_value=None,
+            ),
+        ):
+            mock_settings.base_dir = MagicMock()
+            mock_path = MagicMock()
+            mock_path.exists.return_value = True
+            mock_settings.base_dir.__truediv__ = MagicMock(return_value=mock_path)
+
+            results = await _track_from_3mf(
+                printer_id=1,
+                archive_id=172,
+                status="completed",
+                print_name="no layer info",
+                handled_trays=handled_trays,
+                printer_manager=printer_manager,
+                db=db,
+                last_layer_num=0,  # also unavailable
+            )
+
+        # 2 segments, equal split: 60g / 2 = 30g each. Last segment uses the
+        # `is_last` remainder branch so it stays at 30.0 too.
+        assert len(results) == 2
+        assert results[0]["weight_used"] == 30.0
+        assert results[1]["weight_used"] == 30.0
+
 
 class TestDecodeMqttMapping:
     """Tests for _decode_mqtt_mapping() — snow-encoded MQTT mapping to global tray IDs."""