mcpforunityserver 8.5.0__py3-none-any.whl → 9.1.0__py3-none-any.whl

transport/legacy/unity_connection.py

@@ -233,14 +233,21 @@ class UnityConnection:
             logger.error(f"Error during receive: {str(e)}")
             raise
 
-    def send_command(self, command_type: str, params: dict[str, Any] = None) -> dict[str, Any]:
-        """Send a command with retry/backoff and port rediscovery. Pings only when requested."""
+    def send_command(self, command_type: str, params: dict[str, Any] = None, max_attempts: int | None = None) -> dict[str, Any]:
+        """Send a command with retry/backoff and port rediscovery. Pings only when requested.
+
+        Args:
+            command_type: The Unity command to send
+            params: Command parameters
+            max_attempts: Maximum retry attempts (None = use config default, 0 = no retries)
+        """
         # Defensive guard: catch empty/placeholder invocations early
         if not command_type:
             raise ValueError("MCP call missing command_type")
         if params is None:
             return MCPResponse(success=False, error="MCP call received with no parameters (client placeholder?)")
-        attempts = max(config.max_retries, 5)
+        attempts = max(config.max_retries,
+                       5) if max_attempts is None else max_attempts
         base_backoff = max(0.5, config.retry_delay)
 
         def read_status_file(target_hash: str | None = None) -> dict | None:
@@ -584,7 +591,7 @@ class UnityConnectionPool:
         raise ConnectionError(
             f"Unity instance '{identifier}' not found. "
             f"Available instances: {available_ids}. "
-            f"Check unity://instances resource for all instances."
+            f"Check mcpforunity://instances resource for all instances."
         )
 
     def get_connection(self, instance_identifier: str | None = None) -> UnityConnection:
@@ -686,28 +693,46 @@ def get_unity_connection(instance_identifier: str | None = None) -> UnityConnect
 # Centralized retry helpers
 # -----------------------------
 
-def _is_reloading_response(resp: object) -> bool:
-    """Return True if the Unity response indicates the editor is reloading.
+def _extract_response_reason(resp: object) -> str | None:
+    """Extract a normalized (lowercase) reason string from a response.
 
-    Supports both raw dict payloads from Unity and MCPResponse objects returned
-    by preflight checks or transport helpers.
+    Returns lowercase reason values to enable case-insensitive comparisons
+    by callers (e.g. _is_reloading_response, refresh_unity).
     """
-    # Structured MCPResponse from preflight/transport
     if isinstance(resp, MCPResponse):
-        # Explicit "please retry" hint from preflight
-        if getattr(resp, "hint", None) == "retry":
-            return True
+        data = getattr(resp, "data", None)
+        if isinstance(data, dict):
+            reason = data.get("reason")
+            if isinstance(reason, str):
+                return reason.lower()
         message_text = f"{resp.message or ''} {resp.error or ''}".lower()
-        return "reload" in message_text
+        if "reload" in message_text:
+            return "reloading"
+        return None
 
-    # Raw Unity payloads
     if isinstance(resp, dict):
         if resp.get("state") == "reloading":
-            return True
+            return "reloading"
+        data = resp.get("data")
+        if isinstance(data, dict):
+            reason = data.get("reason")
+            if isinstance(reason, str):
+                return reason.lower()
         message_text = (resp.get("message") or resp.get("error") or "").lower()
-        return "reload" in message_text
+        if "reload" in message_text:
+            return "reloading"
+        return None
 
-    return False
+    return None
+
+
+def _is_reloading_response(resp: object) -> bool:
+    """Return True if the Unity response indicates the editor is reloading.
+
+    Supports both raw dict payloads from Unity and MCPResponse objects returned
+    by preflight checks or transport helpers.
+    """
+    return _extract_response_reason(resp) == "reloading"
 
 
 def send_command_with_retry(
@@ -716,7 +741,8 @@ def send_command_with_retry(
     *,
     instance_id: str | None = None,
     max_retries: int | None = None,
-    retry_ms: int | None = None
+    retry_ms: int | None = None,
+    retry_on_reload: bool = True
 ) -> dict[str, Any] | MCPResponse:
     """Send a command to a Unity instance, waiting politely through Unity reloads.
 
@@ -726,6 +752,8 @@ def send_command_with_retry(
         instance_id: Optional Unity instance identifier (name, hash, name@hash, etc.)
         max_retries: Maximum number of retries for reload states
         retry_ms: Delay between retries in milliseconds
+        retry_on_reload: If False, don't retry when Unity is reloading (for commands
+            that trigger compilation/reload and shouldn't be re-sent)
 
     Returns:
         Response dictionary or MCPResponse from Unity
@@ -738,15 +766,87 @@ def send_command_with_retry(
         max_retries = getattr(config, "reload_max_retries", 40)
     if retry_ms is None:
         retry_ms = getattr(config, "reload_retry_ms", 250)
-
-    response = conn.send_command(command_type, params)
+    try:
+        max_wait_s = float(os.environ.get(
+            "UNITY_MCP_RELOAD_MAX_WAIT_S", "2.0"))
+    except ValueError as e:
+        raw_val = os.environ.get("UNITY_MCP_RELOAD_MAX_WAIT_S", "2.0")
+        logger.warning(
+            "Invalid UNITY_MCP_RELOAD_MAX_WAIT_S=%r, using default 2.0: %s",
+            raw_val, e)
+        max_wait_s = 2.0
+    # Clamp to [0, 30] to prevent misconfiguration from causing excessive waits
+    max_wait_s = max(0.0, min(max_wait_s, 30.0))
+
+    # If retry_on_reload=False, disable connection-level retries too (issue #577)
+    # Commands that trigger compilation/reload shouldn't retry on disconnect
+    send_max_attempts = None if retry_on_reload else 0
+
+    response = conn.send_command(
+        command_type, params, max_attempts=send_max_attempts)
     retries = 0
-    while _is_reloading_response(response) and retries < max_retries:
-        delay_ms = int(response.get("retry_after_ms", retry_ms)
-                       ) if isinstance(response, dict) else retry_ms
-        time.sleep(max(0.0, delay_ms / 1000.0))
+    wait_started = None
+    reason = _extract_response_reason(response)
+    while retry_on_reload and _is_reloading_response(response) and retries < max_retries:
+        if wait_started is None:
+            wait_started = time.monotonic()
+            logger.debug(
+                "Unity reload wait started: command=%s instance=%s reason=%s max_wait_s=%.2f",
+                command_type,
+                instance_id or "default",
+                reason or "reloading",
+                max_wait_s,
+            )
+        if max_wait_s <= 0:
+            break
+        elapsed = time.monotonic() - wait_started
+        if elapsed >= max_wait_s:
+            break
+        delay_ms = retry_ms
+        if isinstance(response, dict):
+            retry_after = response.get("retry_after_ms")
+            if retry_after is None and isinstance(response.get("data"), dict):
+                retry_after = response["data"].get("retry_after_ms")
+            if retry_after is not None:
+                delay_ms = int(retry_after)
+        sleep_ms = max(50, min(int(delay_ms), 250))
+        logger.debug(
+            "Unity reload wait retry: command=%s instance=%s reason=%s retry_after_ms=%s sleep_ms=%s",
+            command_type,
+            instance_id or "default",
+            reason or "reloading",
+            delay_ms,
+            sleep_ms,
+        )
+        time.sleep(max(0.0, sleep_ms / 1000.0))
         retries += 1
         response = conn.send_command(command_type, params)
+        reason = _extract_response_reason(response)
+
+    if wait_started is not None:
+        waited = time.monotonic() - wait_started
+        if _is_reloading_response(response):
+            logger.debug(
+                "Unity reload wait exceeded budget: command=%s instance=%s waited_s=%.3f",
+                command_type,
+                instance_id or "default",
+                waited,
+            )
+            return MCPResponse(
+                success=False,
+                error="Unity is reloading; please retry",
+                hint="retry",
+                data={
+                    "reason": "reloading",
+                    "retry_after_ms": min(250, max(50, retry_ms)),
+                },
+            )
+        logger.debug(
+            "Unity reload wait completed: command=%s instance=%s waited_s=%.3f",
+            command_type,
+            instance_id or "default",
+            waited,
+        )
     return response
 
 
@@ -757,7 +857,8 @@ async def async_send_command_with_retry(
     instance_id: str | None = None,
     loop=None,
     max_retries: int | None = None,
-    retry_ms: int | None = None
+    retry_ms: int | None = None,
+    retry_on_reload: bool = True
 ) -> dict[str, Any] | MCPResponse:
     """Async wrapper that runs the blocking retry helper in a thread pool.
 
@@ -768,6 +869,7 @@ async def async_send_command_with_retry(
         loop: Optional asyncio event loop
         max_retries: Maximum number of retries for reload states
         retry_ms: Delay between retries in milliseconds
+        retry_on_reload: If False, don't retry when Unity is reloading
 
     Returns:
         Response dictionary or MCPResponse on error
@@ -779,7 +881,8 @@ async def async_send_command_with_retry(
         return await loop.run_in_executor(
             None,
             lambda: send_command_with_retry(
-                command_type, params, instance_id=instance_id, max_retries=max_retries, retry_ms=retry_ms),
+                command_type, params, instance_id=instance_id, max_retries=max_retries,
+                retry_ms=retry_ms, retry_on_reload=retry_on_reload),
         )
     except Exception as e:
         return MCPResponse(success=False, error=str(e))

transport/plugin_hub.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import os
 import time
 import uuid
 from typing import Any
@@ -12,6 +13,7 @@ from starlette.endpoints import WebSocketEndpoint
 from starlette.websockets import WebSocket
 
 from core.config import config
+from models.models import MCPResponse
 from transport.plugin_registry import PluginRegistry
 from transport.models import (
     WelcomeMessage,
@@ -28,6 +30,14 @@ from transport.models import (
 logger = logging.getLogger("mcp-for-unity-server")
 
 
+class PluginDisconnectedError(RuntimeError):
+    """Raised when a plugin WebSocket disconnects while commands are in flight."""
+
+
+class NoUnitySessionError(RuntimeError):
+    """Raised when no Unity plugins are available."""
+
+
 class PluginHub(WebSocketEndpoint):
     """Manages persistent WebSocket connections to Unity plugins."""
 
@@ -35,10 +45,19 @@ class PluginHub(WebSocketEndpoint):
     KEEP_ALIVE_INTERVAL = 15
     SERVER_TIMEOUT = 30
     COMMAND_TIMEOUT = 30
+    # Timeout (seconds) for fast-fail commands like ping/read_console/get_editor_state.
+    # Keep short so MCP clients aren't blocked during Unity compilation/reload/unfocused throttling.
+    FAST_FAIL_TIMEOUT = 2.0
+    # Fast-path commands should never block the client for long; return a retry hint instead.
+    # This helps avoid the Cursor-side ~30s tool-call timeout when Unity is compiling/reloading
+    # or is throttled while unfocused.
+    _FAST_FAIL_COMMANDS: set[str] = {
+        "read_console", "get_editor_state", "ping"}
 
     _registry: PluginRegistry | None = None
     _connections: dict[str, WebSocket] = {}
-    _pending: dict[str, asyncio.Future] = {}
+    # command_id -> {"future": Future, "session_id": str}
+    _pending: dict[str, dict[str, Any]] = {}
     _lock: asyncio.Lock | None = None
     _loop: asyncio.AbstractEventLoop | None = None
 
@@ -95,6 +114,22 @@ class PluginHub(WebSocketEndpoint):
                 (sid for sid, ws in cls._connections.items() if ws is websocket), None)
             if session_id:
                 cls._connections.pop(session_id, None)
+                # Fail-fast any in-flight commands for this session to avoid waiting for COMMAND_TIMEOUT.
+                pending_ids = [
+                    command_id
+                    for command_id, entry in cls._pending.items()
+                    if entry.get("session_id") == session_id
+                ]
+                for command_id in pending_ids:
+                    entry = cls._pending.get(command_id)
+                    future = entry.get("future") if isinstance(
+                        entry, dict) else None
+                    if future and not future.done():
+                        future.set_exception(
+                            PluginDisconnectedError(
+                                f"Unity plugin session {session_id} disconnected while awaiting command_result"
+                            )
+                        )
                 if cls._registry:
                     await cls._registry.unregister(session_id)
                 logger.info(
@@ -108,6 +143,36 @@ class PluginHub(WebSocketEndpoint):
         websocket = await cls._get_connection(session_id)
         command_id = str(uuid.uuid4())
         future: asyncio.Future = asyncio.get_running_loop().create_future()
+        # Compute a per-command timeout:
+        # - fast-path commands: short timeout (encourage retry)
+        # - long-running commands: allow caller to request a longer timeout via params
+        unity_timeout_s = float(cls.COMMAND_TIMEOUT)
+        server_wait_s = float(cls.COMMAND_TIMEOUT)
+        if command_type in cls._FAST_FAIL_COMMANDS:
+            fast_timeout = float(cls.FAST_FAIL_TIMEOUT)
+            unity_timeout_s = fast_timeout
+            server_wait_s = fast_timeout
+        else:
+            # Common tools pass a requested timeout in seconds (e.g., timeout_seconds=900).
+            requested = None
+            try:
+                if isinstance(params, dict):
+                    requested = params.get("timeout_seconds", None)
+                    if requested is None:
+                        requested = params.get("timeoutSeconds", None)
+            except Exception:
+                requested = None
+
+            if requested is not None:
+                try:
+                    requested_s = float(requested)
+                    # Clamp to a sane upper bound to avoid accidental infinite hangs.
+                    requested_s = max(1.0, min(requested_s, 60.0 * 60.0))
+                    unity_timeout_s = max(unity_timeout_s, requested_s)
+                    # Give the server a small cushion beyond the Unity-side timeout to account for transport overhead.
+                    server_wait_s = max(server_wait_s, requested_s + 5.0)
+                except Exception:
+                    pass
 
         lock = cls._lock
         if lock is None:
@@ -117,18 +182,36 @@ class PluginHub(WebSocketEndpoint):
             if command_id in cls._pending:
                 raise RuntimeError(
                     f"Duplicate command id generated: {command_id}")
-            cls._pending[command_id] = future
+            cls._pending[command_id] = {
+                "future": future, "session_id": session_id}
 
         try:
             msg = ExecuteCommandMessage(
                 id=command_id,
                 name=command_type,
                 params=params,
-                timeout=cls.COMMAND_TIMEOUT,
+                timeout=unity_timeout_s,
             )
-            await websocket.send_json(msg.model_dump())
-            result = await asyncio.wait_for(future, timeout=cls.COMMAND_TIMEOUT)
-            return result
+            try:
+                await websocket.send_json(msg.model_dump())
+            except Exception as exc:
+                # If send fails (socket already closing), fail the future so callers don't hang.
+                if not future.done():
+                    future.set_exception(exc)
+                raise
+            try:
+                result = await asyncio.wait_for(future, timeout=server_wait_s)
+                return result
+            except PluginDisconnectedError as exc:
+                return MCPResponse(success=False, error=str(exc), hint="retry").model_dump()
+            except asyncio.TimeoutError:
+                if command_type in cls._FAST_FAIL_COMMANDS:
+                    return MCPResponse(
+                        success=False,
+                        error=f"Unity did not respond to '{command_type}' within {server_wait_s:.1f}s; please retry",
+                        hint="retry",
+                    ).model_dump()
+                raise
         finally:
             async with lock:
                 cls._pending.pop(command_id, None)
@@ -232,6 +315,23 @@ class PluginHub(WebSocketEndpoint):
         logger.info(
             f"Registered {len(payload.tools)} tools for session {session_id}")
 
+        try:
+            from services.custom_tool_service import CustomToolService
+
+            service = CustomToolService.get_instance()
+            service.register_global_tools(payload.tools)
+        except RuntimeError as exc:
+            logger.debug(
+                "Skipping global custom tool registration: CustomToolService not initialized yet (%s)",
+                exc,
+            )
+        except Exception as exc:
+            logger.warning(
+                "Unexpected error during global custom tool registration; "
+                "custom tools may not be available globally",
+                exc_info=exc,
+            )
+
     async def _handle_command_result(self, payload: CommandResultMessage) -> None:
         cls = type(self)
         lock = cls._lock
@@ -245,7 +345,8 @@ class PluginHub(WebSocketEndpoint):
             return
 
         async with lock:
-            future = cls._pending.get(command_id)
+            entry = cls._pending.get(command_id)
+        future = entry.get("future") if isinstance(entry, dict) else None
         if future and not future.done():
             future.set_result(result)
 
@@ -284,11 +385,21 @@ class PluginHub(WebSocketEndpoint):
         if cls._registry is None:
             raise RuntimeError("Plugin registry not configured")
 
-        # Use the same defaults as the stdio transport reload handling so that
-        # HTTP/WebSocket and TCP behave consistently without per-project env.
-        max_retries = max(1, int(getattr(config, "reload_max_retries", 40)))
+        # Bound waiting for Unity sessions so calls fail fast when editors are not ready.
+        try:
+            max_wait_s = float(
+                os.environ.get("UNITY_MCP_SESSION_RESOLVE_MAX_WAIT_S", "2.0"))
+        except ValueError as e:
+            raw_val = os.environ.get(
+                "UNITY_MCP_SESSION_RESOLVE_MAX_WAIT_S", "2.0")
+            logger.warning(
+                "Invalid UNITY_MCP_SESSION_RESOLVE_MAX_WAIT_S=%r, using default 2.0: %s",
+                raw_val, e)
+            max_wait_s = 2.0
+        # Clamp to [0, 30] to prevent misconfiguration from causing excessive waits
+        max_wait_s = max(0.0, min(max_wait_s, 30.0))
         retry_ms = float(getattr(config, "reload_retry_ms", 250))
-        sleep_seconds = max(0.05, retry_ms / 1000.0)
+        sleep_seconds = max(0.05, min(0.25, retry_ms / 1000.0))
 
         # Allow callers to provide either just the hash or Name@hash
         target_hash: str | None = None
@@ -317,7 +428,7 @@ class PluginHub(WebSocketEndpoint):
             return None, count
 
         session_id, session_count = await _try_once()
-        deadline = time.monotonic() + (max_retries * sleep_seconds)
+        deadline = time.monotonic() + max_wait_s
         wait_started = None
 
         # If there is no active plugin yet (e.g., Unity starting up or reloading),
@@ -326,33 +437,40 @@ class PluginHub(WebSocketEndpoint):
             if not target_hash and session_count > 1:
                 raise RuntimeError(
                     "Multiple Unity instances are connected. "
-                    "Call set_active_instance with Name@hash from unity://instances."
+                    "Call set_active_instance with Name@hash from mcpforunity://instances."
                 )
             if wait_started is None:
                 wait_started = time.monotonic()
                 logger.debug(
-                    f"No plugin session available (instance={unity_instance or 'default'}); waiting up to {deadline - wait_started:.2f}s",
+                    "No plugin session available (instance=%s); waiting up to %.2fs",
+                    unity_instance or "default",
+                    max_wait_s,
                 )
             await asyncio.sleep(sleep_seconds)
             session_id, session_count = await _try_once()
 
         if session_id is not None and wait_started is not None:
             logger.debug(
-                f"Plugin session restored after {time.monotonic() - wait_started:.3f}s (instance={unity_instance or 'default'})",
+                "Plugin session restored after %.3fs (instance=%s)",
+                time.monotonic() - wait_started,
+                unity_instance or "default",
             )
         if session_id is None and not target_hash and session_count > 1:
             raise RuntimeError(
                 "Multiple Unity instances are connected. "
-                "Call set_active_instance with Name@hash from unity://instances."
+                "Call set_active_instance with Name@hash from mcpforunity://instances."
             )
 
         if session_id is None:
             logger.warning(
-                f"No Unity plugin reconnected within {max_retries * sleep_seconds:.2f}s (instance={unity_instance or 'default'})",
+                "No Unity plugin reconnected within %.2fs (instance=%s)",
+                max_wait_s,
+                unity_instance or "default",
             )
             # At this point we've given the plugin ample time to reconnect; surface
             # a clear error so the client can prompt the user to open Unity.
-            raise RuntimeError("No Unity plugins are currently connected")
+            raise NoUnitySessionError(
+                "No Unity plugins are currently connected")
 
         return session_id
 
@@ -363,7 +481,61 @@ class PluginHub(WebSocketEndpoint):
         command_type: str,
         params: dict[str, Any],
     ) -> dict[str, Any]:
-        session_id = await cls._resolve_session_id(unity_instance)
+        try:
+            session_id = await cls._resolve_session_id(unity_instance)
+        except NoUnitySessionError:
+            logger.debug(
+                "Unity session unavailable; returning retry: command=%s instance=%s",
+                command_type,
+                unity_instance or "default",
+            )
+            return MCPResponse(
+                success=False,
+                error="Unity session not available; please retry",
+                hint="retry",
+                data={"reason": "no_unity_session", "retry_after_ms": 250},
+            ).model_dump()
+
+        # During domain reload / immediate reconnect windows, the plugin may be connected but not yet
+        # ready to process execute commands on the Unity main thread (which can be further delayed when
+        # the Unity Editor is unfocused). For fast-path commands, we do a bounded readiness probe using
+        # a main-thread ping command (handled by TransportCommandDispatcher) rather than waiting on
+        # register_tools (which can be delayed by EditorApplication.delayCall).
+        if command_type in cls._FAST_FAIL_COMMANDS and command_type != "ping":
+            try:
+                max_wait_s = float(os.environ.get(
+                    "UNITY_MCP_SESSION_READY_WAIT_SECONDS", "6"))
+            except ValueError as e:
+                raw_val = os.environ.get(
+                    "UNITY_MCP_SESSION_READY_WAIT_SECONDS", "6")
+                logger.warning(
+                    "Invalid UNITY_MCP_SESSION_READY_WAIT_SECONDS=%r, using default 6.0: %s",
+                    raw_val, e)
+                max_wait_s = 6.0
+            max_wait_s = max(0.0, min(max_wait_s, 30.0))
+            if max_wait_s > 0:
+                deadline = time.monotonic() + max_wait_s
+                while time.monotonic() < deadline:
+                    try:
+                        probe = await cls.send_command(session_id, "ping", {})
+                    except Exception:
+                        probe = None
+
+                    # The Unity-side dispatcher responds with {status:"success", result:{message:"pong"}}
+                    if isinstance(probe, dict) and probe.get("status") == "success":
+                        result = probe.get("result") if isinstance(
+                            probe.get("result"), dict) else {}
+                        if result.get("message") == "pong":
+                            break
+                    await asyncio.sleep(0.1)
+                else:
+                    # Not ready within the bounded window: return retry hint without sending.
+                    return MCPResponse(
+                        success=False,
+                        error=f"Unity session not ready for '{command_type}' (ping not answered); please retry",
+                        hint="retry",
+                    ).model_dump()
+
         return await cls.send_command(session_id, command_type, params)
 
     # ------------------------------------------------------------------