Fix critical issues identified in PR review

This commit addresses all critical and important issues found during
the comprehensive PR review by three specialized agents.

## Critical Fixes

1. **Fixed RobotState immutability conflict** (multi_robot_coordinator.py)
   - Removed numpy array immutability that conflicted with update_robot_state()
   - Arrays are now mutable to allow state updates as intended
   - Updated docstring to clarify arrays remain mutable

2. **Fixed CoordinatedTask field definition** (multi_robot_coordinator.py)
   - Moved completion_callback field before __post_init__ method
   - Was incorrectly defined after __post_init__, causing syntax error

3. **Fixed state extraction in RL environment** (rl_integration.py)
   - Changed from response.get("state", {}).get("qpos") to response.get("qpos")
   - Server returns qpos/qvel directly in response, not nested under "state"
   - Added comment explaining the correct structure

4. **Added missing logger** (rl_integration.py)
   - Added module-level logger = logging.getLogger(__name__)
   - Fixes NameError when logger.error() was called at line 692

5. **Added error handling to ModelViewer.__init__** (mujoco_viewer_server.py)
   - Wrapped model loading in try/except with specific error types
   - Added context-rich error messages for debugging
   - Handles FileNotFoundError, generic model loading errors separately
   - Added error handling for MjData creation and viewer launch

6. **Replaced dangerous BaseException suppression** (mujoco_viewer_server.py)
   - Replaced contextlib.suppress(BaseException) with specific exception types
   - Never suppresses KeyboardInterrupt or SystemExit
   - Added proper error logging for cleanup failures
   - Thread timeout warning when simulation thread doesn't terminate

## Important Fixes

7. **Added thread safety to _handle_ping** (mujoco_viewer_server.py)
   - Acquire viewer_lock before accessing current_viewer and current_model_id
   - Prevents race conditions with concurrent model loading

8. **Improved exception handling in handle_command** (mujoco_viewer_server.py)
   - Distinguish between expected errors (KeyError, ValueError, TypeError)
   - Handle RuntimeError separately (expected runtime failures)
   - Log unexpected exceptions with full stack traces
   - Better error messages for users vs. bugs

9. **Fixed connection state in viewer_client** (viewer_client.py)
   - Mark connection as failed for JSONDecodeError and UnicodeDecodeError
   - Previously only OSError marked connection as failed
   - Prevents continued attempts to use corrupted connections
   - Updated docstring to reflect ValueError instead of JSONDecodeError

## Impact

These fixes address:
- 2 critical bugs that would cause runtime failures
- 1 syntax error in dataclass definition
- 1 missing import causing NameError
- 2 dangerous exception handling patterns
- 3 thread safety and error handling improvements

All changes preserve functionality while significantly improving:
- Error handling robustness
- Thread safety
- Error message clarity
- Connection state consistency

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: robotlearning123

mujoco_viewer_server.py

@@ -46,18 +46,34 @@ def __init__(self, model_id: str, model_source: str):
         self.created_time = time.time()
 
         # Load model - supports file path or XML string
-        if os.path.exists(model_source):
-            # If it's a file path, use from_xml_path to load
-            # (so relative paths are resolved correctly)
-            self.model = mujoco.MjModel.from_xml_path(model_source)
-        else:
-            # Otherwise assume it's an XML string
-            self.model = mujoco.MjModel.from_xml_string(model_source)
+        # Paths are resolved relative to the XML file's directory
+        try:
+            if os.path.exists(model_source):
+                logger.info(f"Loading model {model_id} from file: {model_source}")
+                self.model = mujoco.MjModel.from_xml_path(model_source)
+            else:
+                logger.info(f"Loading model {model_id} from XML string")
+                self.model = mujoco.MjModel.from_xml_string(model_source)
+        except FileNotFoundError as e:
+            logger.error(f"Model file not found for {model_id}: {model_source}")
+            raise RuntimeError(f"Failed to load model {model_id}: file not found at {model_source}") from e
+        except Exception as e:
+            logger.error(f"Failed to load MuJoCo model {model_id}: {e}")
+            raise RuntimeError(f"Failed to load model {model_id}: {e}") from e
 
-        self.data = mujoco.MjData(self.model)
+        # Create simulation data
+        try:
+            self.data = mujoco.MjData(self.model)
+        except Exception as e:
+            logger.error(f"Failed to create MjData for model {model_id}: {e}")
+            raise RuntimeError(f"Failed to initialize simulation data for {model_id}: {e}") from e
 
         # Start viewer
-        self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
+        try:
+            self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
+        except Exception as e:
+            logger.error(f"Failed to launch viewer for model {model_id}: {e}")
+            raise RuntimeError(f"Failed to start viewer for {model_id}: {e}") from e
 
         # Start simulation loop
         self.simulation_running = True
@@ -107,13 +123,25 @@ def close(self):
                     self.viewer.close()
                 elif hasattr(self.viewer, "_window") and self.viewer._window:
                     # For older MuJoCo versions, try to close the window directly
-                    with contextlib.suppress(builtins.BaseException):
+                    try:
                         self.viewer._window.close()
+                    except (AttributeError, RuntimeError) as e:
+                        logger.debug(f"Failed to close viewer window for {self.model_id}: {e}")
+
                 # Wait for simulation thread to finish
                 if hasattr(self, "sim_thread") and self.sim_thread.is_alive():
                     self.sim_thread.join(timeout=2.0)
-            except Exception as e:
+                    if self.sim_thread.is_alive():
+                        logger.warning(f"Simulation thread for {self.model_id} did not terminate within timeout")
+            except KeyboardInterrupt:
+                # Never suppress user interrupts
+                raise
+            except (AttributeError, RuntimeError, OSError) as e:
+                # Expected errors during cleanup
                 logger.warning(f"Error closing viewer for {self.model_id}: {e}")
+            except Exception as e:
+                # Unexpected errors should be logged as errors
+                logger.error(f"Unexpected error closing viewer for {self.model_id}: {e}")
             finally:
                 self.viewer = None
         logger.info(f"Closed ModelViewer for {self.model_id}")
@@ -160,11 +188,21 @@ def handle_command(self, command: Dict[str, Any]) -> Dict[str, Any]:
             handler = self._command_handlers.get(cmd_type)
             if handler:
                 return handler(command)
+            logger.warning(f"Unknown command type received: {cmd_type}")
             return {"success": False, "error": f"Unknown command: {cmd_type}"}
 
-        except Exception as e:
-            logger.exception(f"Error handling command {cmd_type}: {e}")
+        except (KeyError, ValueError, TypeError) as e:
+            # Expected parameter validation errors
+            logger.warning(f"Invalid parameters for command {cmd_type}: {e}")
+            return {"success": False, "error": f"Invalid parameters: {e}"}
+        except RuntimeError as e:
+            # Expected runtime errors (model loading failures, etc.)
+            logger.error(f"Runtime error handling command {cmd_type}: {e}")
             return {"success": False, "error": str(e)}
+        except Exception as e:
+            # Unexpected errors - these indicate bugs
+            logger.exception(f"Unexpected error handling command {cmd_type}: {e}")
+            return {"success": False, "error": f"Internal server error: {str(e)}"}
 
     def _check_viewer_available(self, model_id: str | None) -> Dict[str, Any] | None:
         """Check if viewer is available for the given model. Returns error dict or None if OK."""
@@ -286,12 +324,14 @@ def _handle_list_models(self, command: Dict[str, Any]) -> Dict[str, Any]:
 
     def _handle_ping(self, command: Dict[str, Any]) -> Dict[str, Any]:
         """Ping the server."""
-        models_count = 1 if self.current_viewer else 0
+        with self.viewer_lock:
+            models_count = 1 if self.current_viewer else 0
+            current_model = self.current_model_id
         return {
             "success": True,
             "pong": True,
             "models_count": models_count,
-            "current_model": self.current_model_id,
+            "current_model": current_model,
             "server_running": self.running,
             "server_info": {
                 "version": "0.7.4",

src/mujoco_mcp/multi_robot_coordinator.py

@@ -59,21 +59,16 @@ class RobotState:
     last_update: float = field(default_factory=time.time)
 
     def __post_init__(self):
-        """Validate robot state dimensions and make arrays immutable."""
+        """Validate robot state dimensions.
+
+        Note: Arrays are kept mutable to allow state updates via update_robot_state().
+        """
         if len(self.joint_positions) != len(self.joint_velocities):
             raise ValueError(
                 f"joint_positions length ({len(self.joint_positions)}) must match "
                 f"joint_velocities length ({len(self.joint_velocities)})"
             )
 
-        # Make numpy arrays immutable
-        self.joint_positions.flags.writeable = False
-        self.joint_velocities.flags.writeable = False
-        if self.end_effector_pos is not None:
-            self.end_effector_pos.flags.writeable = False
-        if self.end_effector_vel is not None:
-            self.end_effector_vel.flags.writeable = False
-
     def is_stale(self, timeout: float = 1.0) -> bool:
         """Check if state is stale"""
         return time.time() - self.last_update > timeout
@@ -91,14 +86,14 @@ class CoordinatedTask:
     timeout: float = 30.0
     status: TaskStatus = TaskStatus.PENDING
     start_time: float | None = None
+    completion_callback: Callable | None = None
 
     def __post_init__(self):
         """Validate coordinated task parameters."""
         if not self.robots:
             raise ValueError("robots list cannot be empty")
         if self.timeout <= 0:
             raise ValueError(f"timeout must be positive, got {self.timeout}")
-    completion_callback: Callable | None = None
 
 
 class CollisionChecker:

src/mujoco_mcp/rl_integration.py

@@ -19,6 +19,8 @@
 from .viewer_client import MuJoCoViewerClient
 from .sensor_feedback import SensorManager
 
+logger = logging.getLogger(__name__)
+
 
 class ActionSpaceType(Enum):
     """Types of action spaces for RL environments."""
@@ -669,9 +671,9 @@ def _get_observation(self) -> np.ndarray:
         response = self.viewer_client.send_command({"type": "get_state", "model_id": self.model_id})
 
         if response.get("success"):
-            state = response.get("state", {})
-            qpos = np.array(state.get("qpos", []))
-            qvel = np.array(state.get("qvel", []))
+            # Extract qpos and qvel directly from response (not nested under "state")
+            qpos = np.array(response.get("qpos", []))
+            qvel = np.array(response.get("qvel", []))
 
             # Combine position and velocity
             observation = np.concatenate([qpos, qvel])

src/mujoco_mcp/viewer_client.py

@@ -93,8 +93,7 @@ def send_command(self, command: Dict[str, Any]) -> Dict[str, Any]:
 
         Raises:
             ConnectionError: If not connected to viewer server.
-            ValueError: If response is too large (>1MB).
-            json.JSONDecodeError: If response is not valid JSON.
+            ValueError: If response is too large (>1MB), not valid JSON, or cannot be decoded as UTF-8.
             OSError: If socket communication fails.
         """
         if not self.connected or not self.socket:
@@ -134,9 +133,11 @@ def send_command(self, command: Dict[str, Any]) -> Dict[str, Any]:
             raise OSError(f"Failed to communicate with viewer server: {e}") from e
         except json.JSONDecodeError as e:
             logger.exception(f"Invalid JSON response: {e}")
-            raise
+            self.connected = False  # Connection is likely corrupted
+            raise ValueError(f"Server returned invalid JSON: {e}") from e
         except UnicodeDecodeError as e:
             logger.exception(f"Response decode error: {e}")
+            self.connected = False  # Connection is likely corrupted
             raise ValueError(f"Failed to decode server response as UTF-8: {e}") from e
 
     def ping(self) -> bool: