Fix CodeRabbit review feedback

- Fix partial framed response handling in port discovery
  Add _recv_exact() helper to ensure complete frame reading
  Prevents healthy Unity instances from being misidentified as offline

- Remove unused default_conn variables in server.py (2 files)
  Fixes Ruff F841 lint error that would block CI/CD

- Preserve sync/async nature of resources in wrapper
  Check if original function is coroutine before wrapping
  Prevents 'dict object is not awaitable' runtime errors

- Fix reconnection to preserve instance_id
  Add instance_id tracking to UnityConnection dataclass
  Reconnection now targets the same Unity instance instead of any available one
  Prevents operations from being applied to wrong project

- Add instance logging to manage_asset for debugging
  Helps troubleshoot multi-instance scenarios

🤖 Generated with Claude Code
Co-Authored-By: Claude <[email protected]>

Author: sakurachan

MCPForUnity/UnityMcpServer~/src/port_discovery.py

@@ -82,15 +82,27 @@ def _try_probe_unity_mcp(port: int) -> bool:
                     s.sendall(header + payload)
 
                     # 3. Receive framed response
-                    response_header = s.recv(8)
-                    if len(response_header) != 8:
+                    # Helper to receive exact number of bytes
+                    def _recv_exact(expected: int) -> bytes | None:
+                        chunks = bytearray()
+                        while len(chunks) < expected:
+                            chunk = s.recv(expected - len(chunks))
+                            if not chunk:
+                                return None
+                            chunks.extend(chunk)
+                        return bytes(chunks)
+
+                    response_header = _recv_exact(8)
+                    if response_header is None:
                         return False
 
                     response_length = struct.unpack('>Q', response_header)[0]
                     if response_length > 10000:  # Sanity check
                         return False
 
-                    response = s.recv(response_length)
+                    response = _recv_exact(response_length)
+                    if response is None:
+                        return False
                     return b'"message":"pong"' in response
                 except Exception as e:
                     logger.debug(f"Port probe failed for {port}: {e}")

MCPForUnity/UnityMcpServer~/src/resources/__init__.py

@@ -21,16 +21,24 @@
 def _create_fixed_wrapper(original_func, has_other_params):
     """
     Factory function to create a wrapper that calls original_func with unity_instance=None.
-    This avoids closure issues in loops.
+    This avoids closure issues in loops and preserves sync/async nature of the original function.
     """
+    is_async = inspect.iscoroutinefunction(original_func)
+
     if has_other_params:
-        # Has other parameters (like mode)
-        async def fixed_wrapper(*args, **kwargs):
-            return await original_func(*args, **kwargs, unity_instance=None)
+        if is_async:
+            async def fixed_wrapper(*args, **kwargs):
+                return await original_func(*args, **kwargs, unity_instance=None)
+        else:
+            def fixed_wrapper(*args, **kwargs):
+                return original_func(*args, **kwargs, unity_instance=None)
     else:
-        # No other parameters, just unity_instance
-        async def fixed_wrapper():
-            return await original_func(unity_instance=None)
+        if is_async:
+            async def fixed_wrapper():
+                return await original_func(unity_instance=None)
+        else:
+            def fixed_wrapper():
+                return original_func(unity_instance=None)
 
     return fixed_wrapper
 

MCPForUnity/UnityMcpServer~/src/server.py

@@ -111,7 +111,7 @@ def _emit_startup():
 
                 # Try to connect to default instance
                 try:
-                    default_conn = _unity_connection_pool.get_connection()
+                    _unity_connection_pool.get_connection()
                     logger.info("Connected to default Unity instance on startup")
 
                     # Record successful Unity connection (deferred)

MCPForUnity/UnityMcpServer~/src/tools/manage_asset.py

@@ -35,7 +35,7 @@ async def manage_asset(
     unity_instance: Annotated[str,
                              "Target Unity instance (project name, hash, or 'Name@hash'). If not specified, uses default instance."] | None = None
 ) -> dict[str, Any]:
-    ctx.info(f"Processing manage_asset: {action}")
+    ctx.info(f"Processing manage_asset: {action} (unity_instance={unity_instance or 'default'})")
     # Coerce 'properties' from JSON string to dict for client compatibility
     if isinstance(properties, str):
         try:

MCPForUnity/UnityMcpServer~/src/unity_connection.py

@@ -38,6 +38,7 @@ class UnityConnection:
     port: int = None  # Will be set dynamically
     sock: socket.socket = None  # Socket for Unity communication
     use_framing: bool = False  # Negotiated per-connection
+    instance_id: str | None = None  # Instance identifier for reconnection
 
     def __post_init__(self):
         """Set port from discovery if not explicitly provided"""
@@ -329,9 +330,24 @@ def read_status_file() -> dict | None:
                 finally:
                     self.sock = None
 
-                # Re-discover port each time
+                # Re-discover the port for this specific instance
                 try:
-                    new_port = PortDiscovery.discover_unity_port()
+                    new_port: int | None = None
+                    if self.instance_id:
+                        # Try to rediscover the specific instance
+                        pool = get_unity_connection_pool()
+                        refreshed = pool.discover_all_instances(force_refresh=True)
+                        match = next((inst for inst in refreshed if inst.id == self.instance_id), None)
+                        if match:
+                            new_port = match.port
+                            logger.debug(f"Rediscovered instance {self.instance_id} on port {new_port}")
+                        else:
+                            logger.warning(f"Instance {self.instance_id} not found during reconnection")
+
+                    # Fallback to generic port discovery if instance-specific discovery failed
+                    if new_port is None:
+                        new_port = PortDiscovery.discover_unity_port()
+
                     if new_port != self.port:
                         logger.info(
                             f"Unity port changed {self.port} -> {new_port}")
@@ -547,14 +563,20 @@ def get_connection(self, instance_identifier: Optional[str] = None) -> UnityConn
         with self._pool_lock:
             if target.id not in self._connections:
                 logger.info(f"Creating new connection to Unity instance: {target.id} (port {target.port})")
-                conn = UnityConnection(port=target.port)
+                conn = UnityConnection(port=target.port, instance_id=target.id)
                 if not conn.connect():
                     raise ConnectionError(
                         f"Failed to connect to Unity instance '{target.id}' on port {target.port}. "
                         f"Ensure the Unity Editor is running."
                     )
                 self._connections[target.id] = conn
             else:
+                # Update existing connection with instance_id and port if changed
+                conn = self._connections[target.id]
+                conn.instance_id = target.id
+                if conn.port != target.port:
+                    logger.info(f"Updating cached port for {target.id}: {conn.port} -> {target.port}")
+                    conn.port = target.port
                 logger.debug(f"Reusing existing connection to: {target.id}")
 
             return self._connections[target.id]

Server/port_discovery.py

@@ -82,15 +82,27 @@ def _try_probe_unity_mcp(port: int) -> bool:
                     s.sendall(header + payload)
 
                     # 3. Receive framed response
-                    response_header = s.recv(8)
-                    if len(response_header) != 8:
+                    # Helper to receive exact number of bytes
+                    def _recv_exact(expected: int) -> bytes | None:
+                        chunks = bytearray()
+                        while len(chunks) < expected:
+                            chunk = s.recv(expected - len(chunks))
+                            if not chunk:
+                                return None
+                            chunks.extend(chunk)
+                        return bytes(chunks)
+
+                    response_header = _recv_exact(8)
+                    if response_header is None:
                         return False
 
                     response_length = struct.unpack('>Q', response_header)[0]
                     if response_length > 10000:  # Sanity check
                         return False
 
-                    response = s.recv(response_length)
+                    response = _recv_exact(response_length)
+                    if response is None:
+                        return False
                     return b'"message":"pong"' in response
                 except Exception as e:
                     logger.debug(f"Port probe failed for {port}: {e}")

Server/resources/__init__.py

@@ -21,16 +21,24 @@
 def _create_fixed_wrapper(original_func, has_other_params):
     """
     Factory function to create a wrapper that calls original_func with unity_instance=None.
-    This avoids closure issues in loops.
+    This avoids closure issues in loops and preserves sync/async nature of the original function.
     """
+    is_async = inspect.iscoroutinefunction(original_func)
+
     if has_other_params:
-        # Has other parameters (like mode)
-        async def fixed_wrapper(*args, **kwargs):
-            return await original_func(*args, **kwargs, unity_instance=None)
+        if is_async:
+            async def fixed_wrapper(*args, **kwargs):
+                return await original_func(*args, **kwargs, unity_instance=None)
+        else:
+            def fixed_wrapper(*args, **kwargs):
+                return original_func(*args, **kwargs, unity_instance=None)
     else:
-        # No other parameters, just unity_instance
-        async def fixed_wrapper():
-            return await original_func(unity_instance=None)
+        if is_async:
+            async def fixed_wrapper():
+                return await original_func(unity_instance=None)
+        else:
+            def fixed_wrapper():
+                return original_func(unity_instance=None)
 
     return fixed_wrapper
 

Server/server.py

@@ -111,7 +111,7 @@ def _emit_startup():
 
                 # Try to connect to default instance
                 try:
-                    default_conn = _unity_connection_pool.get_connection()
+                    _unity_connection_pool.get_connection()
                     logger.info("Connected to default Unity instance on startup")
 
                     # Record successful Unity connection (deferred)

Server/tools/manage_asset.py

@@ -35,7 +35,7 @@ async def manage_asset(
     unity_instance: Annotated[str,
                              "Target Unity instance (project name, hash, or 'Name@hash'). If not specified, uses default instance."] | None = None
 ) -> dict[str, Any]:
-    ctx.info(f"Processing manage_asset: {action}")
+    ctx.info(f"Processing manage_asset: {action} (unity_instance={unity_instance or 'default'})")
     # Coerce 'properties' from JSON string to dict for client compatibility
     if isinstance(properties, str):
         try:

Server/unity_connection.py

@@ -38,6 +38,7 @@ class UnityConnection:
     port: int = None  # Will be set dynamically
     sock: socket.socket = None  # Socket for Unity communication
     use_framing: bool = False  # Negotiated per-connection
+    instance_id: str | None = None  # Instance identifier for reconnection
 
     def __post_init__(self):
         """Set port from discovery if not explicitly provided"""
@@ -329,9 +330,24 @@ def read_status_file() -> dict | None:
                 finally:
                     self.sock = None
 
-                # Re-discover port each time
+                # Re-discover the port for this specific instance
                 try:
-                    new_port = PortDiscovery.discover_unity_port()
+                    new_port: int | None = None
+                    if self.instance_id:
+                        # Try to rediscover the specific instance
+                        pool = get_unity_connection_pool()
+                        refreshed = pool.discover_all_instances(force_refresh=True)
+                        match = next((inst for inst in refreshed if inst.id == self.instance_id), None)
+                        if match:
+                            new_port = match.port
+                            logger.debug(f"Rediscovered instance {self.instance_id} on port {new_port}")
+                        else:
+                            logger.warning(f"Instance {self.instance_id} not found during reconnection")
+
+                    # Fallback to generic port discovery if instance-specific discovery failed
+                    if new_port is None:
+                        new_port = PortDiscovery.discover_unity_port()
+
                     if new_port != self.port:
                         logger.info(
                             f"Unity port changed {self.port} -> {new_port}")
@@ -547,14 +563,20 @@ def get_connection(self, instance_identifier: Optional[str] = None) -> UnityConn
         with self._pool_lock:
             if target.id not in self._connections:
                 logger.info(f"Creating new connection to Unity instance: {target.id} (port {target.port})")
-                conn = UnityConnection(port=target.port)
+                conn = UnityConnection(port=target.port, instance_id=target.id)
                 if not conn.connect():
                     raise ConnectionError(
                         f"Failed to connect to Unity instance '{target.id}' on port {target.port}. "
                         f"Ensure the Unity Editor is running."
                     )
                 self._connections[target.id] = conn
             else:
+                # Update existing connection with instance_id and port if changed
+                conn = self._connections[target.id]
+                conn.instance_id = target.id
+                if conn.port != target.port:
+                    logger.info(f"Updating cached port for {target.id}: {conn.port} -> {target.port}")
+                    conn.port = target.port
                 logger.debug(f"Reusing existing connection to: {target.id}")
 
             return self._connections[target.id]