fix: macOS path validation + eject_robot state preservation + mesh init

Three bugs fixed in one commit:

1. **macOS path validation** (_path_validation.py)
   validate_save_path() was silently accepting /etc/passwd and other
   sensitive paths on macOS because os.path.realpath resolves /etc,
   /var, /tmp to /private/etc, /private/var, /private/tmp via symlinks.
   The blocked-prefix startswith check then missed them.

   Fix: on darwin, add /private/-prefixed variants of every Linux
   blocked prefix to BLOCKED_PREFIXES.

   Tests: 5 previously-failing tests now pass, plus 2 new regression
   tests (test_darwin_includes_private_variants,
   test_linux_excludes_private_variants) pin the platform semantics.

2. **eject_robot state preservation** (scene_ops.py)
   eject_robot_from_scene rebuilt the scene from scratch, silently
   resetting surviving robots' joints to qpos=0 and snapping objects
   back to their spawn pose. Agents calling remove_robot mid-scene
   lost physics state with no warning.

   Fix: snapshot per-joint (qpos, qvel) by fully-qualified name BEFORE
   rebuild, restore by name AFTER fresh compile. Matches the 'Per-name
   state copy, not flat index' rule from AGENTS.md — flat-index copy
   is unsafe because body/joint indices shift when a robot is removed.

   New helpers: _snapshot_joint_state, _restore_joint_state (both
   module-private, name-based, width-checked per joint type).

   Tests: 3 new regression tests
   (test_surviving_robot_joint_state_is_preserved,
   test_surviving_object_freejoint_pose_is_preserved,
   test_ejected_robot_state_is_not_restored).

3. **mesh / peer_id explicit init** (simulation.py) — addresses
   @sundargthb's approval nit.

   mesh and peer_id constructor params were never assigned to self,
   making hasattr(self, 'mesh') at L2006 permanently False and
   masking the future PR #98 wire-up. Now explicitly initialised as
   attributes so the guard becomes a plain truthy check.

   Future-compatible: when PR #98 lands and replaces self.mesh with a
   real mesh object, the cleanup path will start working without
   further changes.

Tests: 1275 pass, 1 skipped, 0 failures (was 1255 pass + 5 fail).
Lint: ruff + mypy clean.

Author: cagataycali

strands_robots/simulation/mujoco/scene_ops.py

@@ -191,6 +191,99 @@ def eject_body_from_scene(world: SimWorld, body_name: str) -> bool:
     return _recompile_preserving_state(world, spec)
 
 
+def _snapshot_joint_state(world: SimWorld) -> dict[str, tuple[list[float], list[float]]]:
+    """Snapshot per-joint ``(qpos, qvel)`` slices keyed by fully-qualified
+    MuJoCo joint name.
+
+    Used by :func:`eject_robot_from_scene` to preserve the state of surviving
+    robots and object freejoints across a scene rebuild. Flat-index slicing
+    is unsafe here because the body-tree order may shift when a robot is
+    removed (see AGENTS.md "Per-name state copy" rule).
+
+    Returns a dict mapping ``<joint_name> -> (qpos_slice, qvel_slice)`` where
+    each slice has the appropriate width for the joint type (1 for hinge/
+    slide, 4 for ball, 7 for free).
+    """
+    if world._model is None or world._data is None:
+        return {}
+    mj = _ensure_mujoco()
+    model = world._model
+    data = world._data
+    snap: dict[str, tuple[list[float], list[float]]] = {}
+    for jid in range(model.njnt):
+        name = mj.mj_id2name(model, mj.mjtObj.mjOBJ_JOINT, jid)
+        if not name:
+            continue
+        qpos_adr = int(model.jnt_qposadr[jid])
+        qvel_adr = int(model.jnt_dofadr[jid])
+        jtype = int(model.jnt_type[jid])
+        # qpos width: free=7, ball=4, hinge/slide=1
+        # qvel width: free=6, ball=3, hinge/slide=1
+        if jtype == mj.mjtJoint.mjJNT_FREE:
+            qpos_w, qvel_w = 7, 6
+        elif jtype == mj.mjtJoint.mjJNT_BALL:
+            qpos_w, qvel_w = 4, 3
+        else:
+            qpos_w, qvel_w = 1, 1
+        snap[name] = (
+            [float(x) for x in data.qpos[qpos_adr : qpos_adr + qpos_w]],
+            [float(x) for x in data.qvel[qvel_adr : qvel_adr + qvel_w]],
+        )
+    return snap
+
+
+def _restore_joint_state(
+    world: SimWorld,
+    snapshot: dict[str, tuple[list[float], list[float]]],
+) -> int:
+    """Restore per-joint state from a snapshot into ``world._data`` by name.
+
+    Joints that no longer exist in the compiled model (e.g. those belonging
+    to the ejected robot) are silently skipped. Joints that exist in the
+    new model but were not in the snapshot keep their fresh-compile defaults
+    (body pos/quat for freejoints, 0 for hinge/slide).
+
+    Returns the number of joints actually restored, for logging.
+    """
+    if world._model is None or world._data is None:
+        return 0
+    mj = _ensure_mujoco()
+    model = world._model
+    data = world._data
+    restored = 0
+    for name, (qpos_vals, qvel_vals) in snapshot.items():
+        jid = mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, name)
+        if jid < 0:
+            continue  # joint no longer exists (expected for ejected robot)
+        qpos_adr = int(model.jnt_qposadr[jid])
+        qvel_adr = int(model.jnt_dofadr[jid])
+        # Width sanity check: if joint type changed (should not happen for
+        # same-name joints across an eject), skip to avoid corrupting state.
+        jtype = int(model.jnt_type[jid])
+        if jtype == mj.mjtJoint.mjJNT_FREE:
+            expect_qp, expect_qv = 7, 6
+        elif jtype == mj.mjtJoint.mjJNT_BALL:
+            expect_qp, expect_qv = 4, 3
+        else:
+            expect_qp, expect_qv = 1, 1
+        if len(qpos_vals) != expect_qp or len(qvel_vals) != expect_qv:
+            logger.warning(
+                "_restore_joint_state: width mismatch for %r (qpos %d!=%d or qvel %d!=%d), skipping",
+                name,
+                len(qpos_vals),
+                expect_qp,
+                len(qvel_vals),
+                expect_qv,
+            )
+            continue
+        for i, v in enumerate(qpos_vals):
+            data.qpos[qpos_adr + i] = v
+        for i, v in enumerate(qvel_vals):
+            data.qvel[qvel_adr + i] = v
+        restored += 1
+    return restored
+
+
 def eject_robot_from_scene(world: SimWorld, robot_name: str) -> bool:
     """Remove every spec element namespaced under ``{robot_name}/``.
 
@@ -199,10 +292,20 @@ def eject_robot_from_scene(world: SimWorld, robot_name: str) -> bool:
     attached child spec's memory gets freed twice). To sidestep that bug
     we REBUILD the scene spec from scratch using the post-remove
     ``world.robots`` / ``world.objects`` / ``world.cameras`` state, then
-    re-attach the remaining robots. Joint state is not preserved across this
-    path - callers that care should call ``reset`` or save/restore state
-    around remove_robot. In the common case (agent removes a robot to clear
-    the scene), this is the expected behaviour anyway.
+    re-attach the remaining robots.
+
+    Joint state preservation: before the rebuild we snapshot every joint's
+    ``(qpos, qvel)`` keyed by fully-qualified name; after the fresh compile
+    we restore state for every joint that still exists in the new model.
+    Joints belonging to the ejected robot are naturally dropped (their name
+    no longer resolves). This keeps surviving robots at their current pose
+    and object freejoints at their current world pose - the behaviour the
+    agent expects when calling ``remove_robot`` mid-scene.
+
+    Flat-index slicing is **not** safe here: removing a robot shifts every
+    body/joint index that comes after it in the kinematic tree, so
+    ``data.qpos[:]`` copies across compiles would mis-assign DOFs. Per-name
+    lookup is the only correct approach (see AGENTS.md).
     """
     spec = _get_spec(world)
     if spec is None or world._model is None:
@@ -211,10 +314,10 @@ def eject_robot_from_scene(world: SimWorld, robot_name: str) -> bool:
 
     mj = _ensure_mujoco()
 
-    # Preserve the current qpos for bodies that are NOT being removed.
-    # We rebuild from world state and then re-attach remaining robots, so
-    # object freejoints start at their body pos (matching fresh add_object
-    # semantics); robot joints start at qpos=0 (same as fresh add_robot).
+    # Snapshot joint state BEFORE we rebuild. Keyed by the fully-qualified
+    # MuJoCo joint name (prefix/joint for attached robots, bare name for
+    # object freejoints).
+    state_snapshot = _snapshot_joint_state(world)
 
     # First drop cameras that originated from the robot being ejected.
     # They're in world.cameras with origin_robot == robot_name. Without this,
@@ -254,6 +357,15 @@ def eject_robot_from_scene(world: SimWorld, robot_name: str) -> bool:
     except Exception as xml_err:
         logger.debug("spec.to_xml() failed: %s", xml_err)
 
+    # Step 4: restore state for every joint that survived the rebuild. Joints
+    # belonging to the ejected robot simply don't resolve and get skipped.
+    restored = _restore_joint_state(world, state_snapshot)
+
+    # Step 5: run a forward pass so derived quantities (xpos, cam xforms)
+    # reflect the restored state. Without this, the next render() call can
+    # produce stale frames because MjData was freshly allocated in Step 3.
+    mj.mj_forward(new_model, new_data)
+
     # Re-discover joint/actuator IDs for remaining robots.
     for robot in world.robots.values():
         pfx = robot.namespace or ""
@@ -274,7 +386,12 @@ def eject_robot_from_scene(world: SimWorld, robot_name: str) -> bool:
         if not robot.actuator_ids and len(world.robots) == 1:
             robot.actuator_ids = list(range(new_model.nu))
 
-    logger.debug("eject_robot %r: scene rebuilt", robot_name)
+    logger.debug(
+        "eject_robot %r: scene rebuilt, restored state for %d/%d joints",
+        robot_name,
+        restored,
+        len(state_snapshot),
+    )
     return True
 
 

strands_robots/simulation/mujoco/simulation.py

@@ -128,12 +128,42 @@ def __init__(
         peer_id: str | None = None,
         **kwargs,
     ):
+        """Construct a MuJoCo Simulation AgentTool.
+
+        Args:
+            tool_name: Identifier surfaced to the agent and used as the
+                thread-name prefix for the executor.
+            default_timestep: Default physics timestep (seconds). Can be
+                overridden via ``create_world(timestep=...)``.
+            default_width: Default render width (pixels) used when a
+                caller does not pass explicit dimensions to ``render``.
+            default_height: Default render height (pixels).
+            mesh: Optional mesh-networking hook. Falsy (default) keeps
+                the Simulation standalone - all mesh code paths are
+                no-ops. When set to a live mesh-client object exposing
+                ``.stop()``, ``cleanup()`` will detach this Simulation
+                from the peer network before tearing down the MuJoCo
+                world. The attribute is plain (not a property), so
+                consumers may attach a client after construction.
+            peer_id: Stable identifier the mesh transport uses to
+                address this Simulation. Opaque to MuJoCo itself; only
+                consulted when ``mesh`` is truthy.
+            **kwargs: Forwarded to ``AgentTool.__init__`` for subclass
+                compatibility.
+        """
         super().__init__()
         self.tool_name_str = tool_name
         self.default_timestep = default_timestep
         self.default_width = default_width
         self.default_height = default_height
 
+        # Mesh attributes are stored plainly (no property wrapper) so
+        # downstream code can swap in a real mesh client after
+        # construction without a setter dance. See the ``mesh`` /
+        # ``peer_id`` docstring entries above for the contract.
+        self.mesh: Any = mesh if mesh else None
+        self.peer_id: str | None = peer_id
+
         self._world: SimWorld | None = None
         self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix=f"{tool_name}_sim")
         # Per-robot Future refs for *active* policies. Completed futures are
@@ -2008,7 +2038,13 @@ def cleanup(self, policy_stop_timeout: float | None = None) -> None:
                 ``_DEFAULT_POLICY_STOP_TIMEOUT`` (5s). Set to a small value
                 in tests that want fast teardown.
         """
-        if hasattr(self, "mesh") and self.mesh:
+        # Detach from the mesh network first (if attached). A truthy
+        # ``self.mesh`` is any object exposing ``.stop()``; falsy values
+        # (the default) mean this Simulation never joined a mesh and
+        # there's nothing to release. Done BEFORE stopping policies so
+        # peer-visible state is torn down cleanly even if the policy
+        # teardown below hits the fallback ``wait=False`` path.
+        if self.mesh:
             self.mesh.stop()
 
         timeout = policy_stop_timeout if policy_stop_timeout is not None else self._DEFAULT_POLICY_STOP_TIMEOUT

strands_robots/tools/_path_validation.py

@@ -44,11 +44,19 @@
 
 
 def _get_blocked_prefixes() -> tuple[str, ...]:
-    """Return blocked prefixes for the current platform."""
+    """Return blocked prefixes for the current platform.
+
+    On macOS, many system directories (``/etc``, ``/var``, ``/tmp``) are
+    symlinks into ``/private/``. Since :func:`validate_save_path` compares
+    against ``os.path.realpath`` output, we must include the ``/private/``-
+    prefixed variants so that ``/etc/passwd`` (which resolves to
+    ``/private/etc/passwd``) is still rejected.
+    """
     if sys.platform == "win32":
         return _WINDOWS_BLOCKED_PREFIXES
     elif sys.platform == "darwin":
-        return _LINUX_BLOCKED_PREFIXES + _MACOS_BLOCKED_PREFIXES
+        private_variants = tuple("/private" + p for p in _LINUX_BLOCKED_PREFIXES)
+        return _LINUX_BLOCKED_PREFIXES + private_variants + _MACOS_BLOCKED_PREFIXES
     else:
         return _LINUX_BLOCKED_PREFIXES
 

tests/simulation/mujoco/test_agenttool_contract.py

@@ -605,3 +605,113 @@ def test_remove_middle_of_three_robots(self, sim):
         assert set(sim._world.robots) == {"alice", "carol"}
         # bob was 6 joints; alice (6) + carol (19) = 25 should remain.
         assert sim._world._model.njnt == njnt_before - 6
+
+
+class TestRemoveRobotPreservesState:
+    """Regression tests for PR #85 follow-up (AGENTS.md "Per-name state
+    copy, not flat index"): removing one robot must NOT reset the state
+    of surviving robots or objects.
+
+    Before the fix, ``eject_robot_from_scene`` rebuilt the scene from
+    scratch and reset every remaining qpos/qvel to 0 (robots) or body
+    pose (objects). Agents that called ``remove_robot`` mid-simulation
+    silently lost physics state.
+    """
+
+    def test_surviving_robot_joint_state_is_preserved(self, sim):
+        """After ``remove_robot(bob)``, alice's joints keep their qpos."""
+        sim.add_robot(name="alice", data_config="so100", position=[-0.5, 0, 0])
+        sim.add_robot(name="bob", data_config="so100", position=[0.5, 0, 0])
+
+        # Drive alice's joints to a non-zero pose and step forward so
+        # qpos actually reflects applied ctrl (not just a zero default).
+        import numpy as np
+
+        alice = sim._world.robots["alice"]
+        # Write directly to qpos for a deterministic snapshot (avoids
+        # ctrl-dynamics dependency). Each joint gets a distinct non-zero
+        # value so accidental index shifts would be obvious.
+        target = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        for jid, val in zip(alice.joint_ids, target):
+            qadr = sim._world._model.jnt_qposadr[jid]
+            sim._world._data.qpos[qadr] = val
+
+        snapshot_before = np.array(
+            [sim._world._data.qpos[sim._world._model.jnt_qposadr[jid]] for jid in alice.joint_ids]
+        )
+
+        # Remove bob - alice should survive with her joint state intact.
+        r = sim.remove_robot(name="bob")
+        assert r["status"] == "success"
+
+        alice_after = sim._world.robots["alice"]
+        snapshot_after = np.array(
+            [sim._world._data.qpos[sim._world._model.jnt_qposadr[jid]] for jid in alice_after.joint_ids]
+        )
+
+        np.testing.assert_allclose(
+            snapshot_before,
+            snapshot_after,
+            atol=1e-10,
+            err_msg="alice's joint qpos was reset by remove_robot(bob)",
+        )
+
+    def test_surviving_object_freejoint_pose_is_preserved(self, sim):
+        """An object's freejoint qpos (position + quat) survives the
+        eject rebuild. Before the fix, objects snapped back to their
+        ``add_object`` spawn pose."""
+        import numpy as np
+
+        sim.add_robot(name="alice", data_config="so100")
+        sim.add_robot(name="bob", data_config="so100", position=[1, 0, 0])
+        sim.add_object(name="cube", shape="box", size=[0.05, 0.05, 0.05], position=[0.3, 0, 0.05])
+
+        mj = sim._mj
+        cube_jid = mj.mj_name2id(sim._world._model, mj.mjtObj.mjOBJ_JOINT, "cube_joint")
+        assert cube_jid >= 0, "cube freejoint must exist"
+        cube_qadr = int(sim._world._model.jnt_qposadr[cube_jid])
+
+        # Move the cube to a distinct position + tilted orientation so
+        # the test fails loudly if state is reset to spawn pose.
+        new_pose = [0.7, -0.2, 0.15, 0.7071, 0.0, 0.7071, 0.0]
+        for i, v in enumerate(new_pose):
+            sim._world._data.qpos[cube_qadr + i] = v
+        mj.mj_forward(sim._world._model, sim._world._data)
+
+        # Eject bob.
+        r = sim.remove_robot(name="bob")
+        assert r["status"] == "success"
+
+        # Cube freejoint must still be there at the moved pose.
+        cube_jid2 = mj.mj_name2id(sim._world._model, mj.mjtObj.mjOBJ_JOINT, "cube_joint")
+        assert cube_jid2 >= 0, "cube freejoint disappeared after remove_robot"
+        cube_qadr2 = int(sim._world._model.jnt_qposadr[cube_jid2])
+        restored = np.array([sim._world._data.qpos[cube_qadr2 + i] for i in range(7)])
+        np.testing.assert_allclose(
+            restored,
+            new_pose,
+            atol=1e-10,
+            err_msg="cube freejoint pose was reset by remove_robot",
+        )
+
+    def test_ejected_robot_state_is_not_restored(self, sim):
+        """The ejected robot's joints should NOT appear in the new model.
+        This is the contrapositive of the preservation tests - confirms
+        the snapshot/restore loop only touches surviving joints.
+        """
+        sim.add_robot(name="alice", data_config="so100")
+        sim.add_robot(name="bob", data_config="so100", position=[1, 0, 0])
+
+        bob = sim._world.robots["bob"]
+        bob_prefix = bob.namespace  # e.g. "bob/"
+        assert bob_prefix, "bob must have a namespace for this test to be meaningful"
+
+        r = sim.remove_robot(name="bob")
+        assert r["status"] == "success"
+
+        mj = sim._mj
+        # No joint under bob's prefix should exist anymore.
+        model = sim._world._model
+        for jid in range(model.njnt):
+            name = mj.mj_id2name(model, mj.mjtObj.mjOBJ_JOINT, jid)
+            assert name is None or not name.startswith(bob_prefix), f"ejected robot's joint survived: {name!r}"

tests/tools/test_path_validation.py

@@ -216,6 +216,26 @@ def test_darwin_includes_macos_extras(self):
             assert "/System/" in prefixes  # macOS specific
             assert "/Library/LaunchDaemons/" in prefixes
 
+    def test_darwin_includes_private_variants(self):
+        """Regression: on macOS, ``os.path.realpath`` maps ``/etc`` →
+        ``/private/etc`` (and same for ``/var``, ``/tmp``). The blocked
+        prefix list MUST include the ``/private/``-prefixed variants,
+        otherwise resolved paths bypass the check entirely. Fixes the
+        bug where ``/etc/passwd`` was silently accepted on macOS."""
+        with patch.object(sys, "platform", "darwin"):
+            prefixes = _get_blocked_prefixes()
+            assert "/private/etc/" in prefixes
+            assert "/private/var/spool/cron/" in prefixes
+            assert "/private/var/spool/at/" in prefixes
+
+    def test_linux_excludes_private_variants(self):
+        """On Linux ``/private/`` is not a special path; the variants
+        should only be added on darwin."""
+        with patch.object(sys, "platform", "linux"):
+            prefixes = _get_blocked_prefixes()
+            for p in prefixes:
+                assert not p.startswith("/private/"), f"Linux should not block /private/* prefixes: {p!r}"
+
     def test_windows_prefixes_returned_on_win32(self):
         """On Windows, Windows-specific prefixes should be active."""
         with patch.object(sys, "platform", "win32"):