From 511544c2eb255557ff3f4a36538e7be84a1d4e17 Mon Sep 17 00:00:00 2001
From: Codex Agent <codex@example.com>
Date: Sat, 22 Nov 2025 18:47:19 +0100
Subject: [PATCH] Persist manual control and pump states

---
 src/reactor_sim/control.py |  2 ++
 src/reactor_sim/reactor.py | 15 +++++++++++++++
 tests/test_simulation.py   |  4 ++++
 3 files changed, 21 insertions(+)

diff --git a/src/reactor_sim/control.py b/src/reactor_sim/control.py
index 536eed6..d201a6e 100644
--- a/src/reactor_sim/control.py
+++ b/src/reactor_sim/control.py
@@ -77,6 +77,7 @@ class ControlSystem:
             "control": {
                 "setpoint_mw": self.setpoint_mw,
                 "rod_fraction": self.rod_fraction,
+                "manual_control": self.manual_control,
             },
             "plant": plant_state.to_dict(),
             "metadata": metadata or {},
@@ -94,6 +95,7 @@ class ControlSystem:
         control = data.get("control", {})
         self.setpoint_mw = control.get("setpoint_mw", self.setpoint_mw)
         self.rod_fraction = control.get("rod_fraction", self.rod_fraction)
+        self.manual_control = control.get("manual_control", self.manual_control)
         plant = PlantState.from_dict(data["plant"])
         LOGGER.info("Loaded plant state from %s", path)
         return plant, data.get("metadata", {}), data.get("health")
diff --git a/src/reactor_sim/reactor.py b/src/reactor_sim/reactor.py
index ddf3ab1..bc46b5e 100644
--- a/src/reactor_sim/reactor.py
+++ b/src/reactor_sim/reactor.py
@@ -461,6 +461,21 @@ class Reactor:
         if health:
             self.health_monitor.load_snapshot(health)
         LOGGER.info("Reactor state restored from %s", filepath)
+        # Back-fill pump state lists for compatibility.
+        if not plant.primary_pumps or len(plant.primary_pumps) < 2:
+            plant.primary_pumps = [
+                PumpState(active=self.primary_pump_active, flow_rate=plant.primary_loop.mass_flow_rate / 2, pressure=plant.primary_loop.pressure)
+                for _ in range(2)
+            ]
+        if not plant.secondary_pumps or len(plant.secondary_pumps) < 2:
+            plant.secondary_pumps = [
+                PumpState(
+                    active=self.secondary_pump_active,
+                    flow_rate=plant.secondary_loop.mass_flow_rate / 2,
+                    pressure=plant.secondary_loop.pressure,
+                )
+                for _ in range(2)
+            ]
         if len(plant.turbines) < len(self.turbines):
             ambient = constants.ENVIRONMENT_TEMPERATURE
             while len(plant.turbines) < len(self.turbines):
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 581d41f..a717f4d 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -20,6 +20,7 @@ def test_reactor_initial_state_is_cold():
 
 def test_state_save_and_load_roundtrip(tmp_path: Path):
     reactor = Reactor.default()
+    reactor.control.manual_control = True
     sim = ReactorSimulation(reactor, timestep=5.0, duration=15.0)
     sim.log()
     save_path = tmp_path / "plant_state.json"
@@ -32,6 +33,9 @@ def test_state_save_and_load_roundtrip(tmp_path: Path):
         sim.last_state.core.fuel_temperature
     )
     assert restored_reactor.control.rod_fraction == reactor.control.rod_fraction
+    assert restored_reactor.control.manual_control == reactor.control.manual_control
+    assert len(restored_state.primary_pumps) == 2
+    assert len(restored_state.secondary_pumps) == 2
 
 
 def test_health_monitor_flags_core_failure():