Jets: enforce stationary rule at ObjectManager.set_body_state too

Move the Jet motion check into ObjectManager.validate_jet_motion (uses the
synced field's cut-link jet_id) and call it from both Simulation.set_body
and the direct ObjectManager.set_body_state entry, before any state or
action change. Test both entries, vx too, and that action is untouched.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/CelerisLab/body/manager.py

@@ -738,6 +738,18 @@ class ObjectManager:
         assert self.obs_pinned is not None
         return self.obs_pinned.copy()
 
+    def validate_jet_motion(self, body_id: int, state: dict) -> None:
+        """Reject motion of a body whose synced cut links are Jet-driven.
+
+        Before the first sync no links exist; sync-time checks cover that case.
+        """
+        field = self._telemetry_field
+        cl = field.curved if field is not None else None
+        if (cl is not None and cl.jet_id is not None
+                and np.any(cl.body_id[cl.jet_id >= 0] == body_id)
+                and any(state.get(k, 0.) != 0 for k in ("vx", "vy", "omega"))):
+            raise ValueError("Jet-driven bodies must stay stationary")
+
     def set_body_state(self, body_id: int, omega: float = 0.0) -> None:
         """Set runtime body state (host action array only, no H2D).
 
@@ -745,7 +757,9 @@ class ObjectManager:
         """
         self._validate_body_id(body_id)
         obj = self.get(body_id)
-        obj.validate_surface_actuation(self.cfg, dict(vars(obj.state), omega=omega))
+        proposed = dict(vars(obj.state), omega=float(omega))
+        obj.validate_surface_actuation(self.cfg, proposed)
+        self.validate_jet_motion(body_id, proposed)
         dim = self.cfg.dim if self.cfg else 2
         slot = 3 * dim
         base = body_id * slot

src/CelerisLab/simulation.py

@@ -260,11 +260,7 @@ class Simulation:
         proposed.update({k: float(v) for k, v in (("omega", omega), ("vx", vx), ("vy", vy))
                          if v is not None})
         obj.validate_surface_actuation(self.lbm_cfg, proposed)
-        cl = self.field.curved
-        if (self._initialized and cl.jet_id is not None
-                and np.any(cl.body_id[cl.jet_id >= 0] == id)
-                and any(proposed.get(k, 0.) != 0 for k in ("vx", "vy", "omega"))):
-            raise ValueError("Jet-driven bodies must stay stationary")
+        self.bodies.validate_jet_motion(id, proposed)
         changed = False
         if omega is not None:
             obj.state.omega = float(omega)

tests/unit/test_jets.py

@@ -243,8 +243,14 @@ def test_set_body_rejects_motion_of_jet_driven_body(device):
     sim = SimpleNamespace(bodies=bodies, field=field, lbm_cfg=field.cfg, _initialized=True)
     with pytest.raises(ValueError, match="stationary"):
         Simulation.set_body(sim, 0, omega=.01)
-    assert bodies.get(0).state.omega == 0
+    with pytest.raises(ValueError, match="stationary"):
+        Simulation.set_body(sim, 0, vx=.01)
+    action = bodies.action.copy()
+    with pytest.raises(ValueError, match="stationary"):
+        bodies.set_body_state(0, omega=.01)  # direct manager entry
+    assert bodies.get(0).state.omega == 0 and np.array_equal(bodies.action, action)
     Simulation.set_body(sim, 0, omega=0.)  # stationary updates stay allowed
+    bodies.set_body_state(0, omega=0.)
 
 
 def test_fine_boundary_keeps_edge_frame_hits_and_checks_jets():