From e18f100e15385bd920500b572aa5ce87a400d099 Mon Sep 17 00:00:00 2001
From: Codex Agent <codex@example.com>
Date: Sat, 22 Nov 2025 19:07:45 +0100
Subject: [PATCH] Handle individual pump unit toggles

---
 src/reactor_sim/dashboard.py | 12 ++++++++++++
 src/reactor_sim/reactor.py   | 26 ++++++++++++++++++++++++++
 tests/test_simulation.py     | 28 ++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/src/reactor_sim/dashboard.py b/src/reactor_sim/dashboard.py
index 2a764e9..361162b 100644
--- a/src/reactor_sim/dashboard.py
+++ b/src/reactor_sim/dashboard.py
@@ -204,6 +204,18 @@ class ReactorDashboard:
         current = self.reactor.turbine_unit_active[index]
         self._queue_command(ReactorCommand(turbine_units={index + 1: not current}))
 
+    def _toggle_primary_pump_unit(self, index: int) -> None:
+        if index < 0 or index >= len(self.reactor.primary_pump_units):
+            return
+        current = self.reactor.primary_pump_units[index]
+        self._queue_command(ReactorCommand(primary_pumps={index + 1: not current}))
+
+    def _toggle_secondary_pump_unit(self, index: int) -> None:
+        if index < 0 or index >= len(self.reactor.secondary_pump_units):
+            return
+        current = self.reactor.secondary_pump_units[index]
+        self._queue_command(ReactorCommand(secondary_pumps={index + 1: not current}))
+
     def _request_reset(self) -> None:
         self.reset_requested = True
         if self.sim:
diff --git a/src/reactor_sim/reactor.py b/src/reactor_sim/reactor.py
index 73499d4..41cdf24 100644
--- a/src/reactor_sim/reactor.py
+++ b/src/reactor_sim/reactor.py
@@ -369,6 +369,32 @@ class Reactor:
         elif active and not any(self.secondary_pump_units):
             self.secondary_pump_units = [True] * len(self.secondary_pump_units)
 
+    def _toggle_primary_pump_unit(self, index: int, active: bool) -> None:
+        if index < 0 or index >= len(self.primary_pump_units):
+            LOGGER.warning("Ignoring primary pump index %s", index)
+            return
+        if self.primary_pump_units[index] == active:
+            return
+        self.primary_pump_units[index] = active
+        LOGGER.info("Primary pump %d %s", index + 1, "enabled" if active else "stopped")
+        if active:
+            self._set_primary_pump(True)
+        elif not any(self.primary_pump_units):
+            self._set_primary_pump(False)
+
+    def _toggle_secondary_pump_unit(self, index: int, active: bool) -> None:
+        if index < 0 or index >= len(self.secondary_pump_units):
+            LOGGER.warning("Ignoring secondary pump index %s", index)
+            return
+        if self.secondary_pump_units[index] == active:
+            return
+        self.secondary_pump_units[index] = active
+        LOGGER.info("Secondary pump %d %s", index + 1, "enabled" if active else "stopped")
+        if active:
+            self._set_secondary_pump(True)
+        elif not any(self.secondary_pump_units):
+            self._set_secondary_pump(False)
+
     def _set_turbine_state(self, active: bool, index: int | None = None) -> None:
         if index is None:
             for idx in range(len(self.turbine_unit_active)):
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index a717f4d..cc3b263 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -101,3 +101,31 @@ def test_toggle_maintenance_progresses_until_restored():
     sim.log()
     assert pump.integrity >= 0.99
     assert "primary_pump" not in reactor.maintenance_active
+
+
+def test_primary_pump_unit_toggle_updates_active_flag():
+    reactor = Reactor.default()
+    state = reactor.initial_state()
+    reactor.primary_pump_active = True
+    reactor.primary_pump_units = [True, True]
+
+    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pumps={1: False}))
+    assert reactor.primary_pump_units == [False, True]
+    assert reactor.primary_pump_active is True
+
+    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pumps={2: False}))
+    assert reactor.primary_pump_units == [False, False]
+    assert reactor.primary_pump_active is False
+
+
+def test_secondary_pump_unit_toggle_can_restart_pump():
+    reactor = Reactor.default()
+    state = reactor.initial_state()
+    reactor.secondary_pump_active = False
+    reactor.secondary_pump_units = [False, False]
+
+    reactor.step(state, dt=1.0, command=ReactorCommand(secondary_pumps={1: True}))
+
+    assert reactor.secondary_pump_units == [True, False]
+    assert reactor.secondary_pump_active is True
+    assert state.secondary_loop.mass_flow_rate > 0.0