diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index fe24bc7725..e5f96dcadd 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -152,6 +152,14 @@ async def execute_tool( message=f"Tool '{tool_name}' not found for project {project_id}", ) + # Custom tools are frequently long-running (captures, image comparison, + # batch imports), so the hub's default 30s command budget times them out + # while they are still working. Give them a generous default; callers can + # still override with an explicit timeout_seconds (or timeoutSeconds) parameter, up to 3600. + # The hub reads timeout_seconds first, so a default added beside a caller's timeoutSeconds would win. + if "timeout_seconds" not in params and "timeoutSeconds" not in params: + params["timeout_seconds"] = 300 + response = await send_with_unity_instance( async_send_command_with_retry, unity_instance, @@ -216,20 +224,45 @@ async def _poll_until_complete( if time.time() > deadline: return MCPResponse( success=False, - message=f"Timeout waiting for {tool_name} to complete", + message=( + f"Timeout waiting for {tool_name} to complete. The operation may " + "still be running in Unity — do not assume it finished or failed; " + "verify its effects before continuing." + ), data=self._safe_response(response), + # No retry hint: the operation may still be running, and resubmitting it could run it twice. ) - await asyncio.sleep(poll_interval) + await asyncio.sleep(min(poll_interval, max(0.0, deadline - time.time()))) + + # A poll inherits the tool's parameters, including its long command budget (300 s by default): + # left as is, one status poll could wait far past this polling deadline. Give each poll what is + # left of the budget instead. The hub still waits at least its 30 s default for any command, + # so the wait itself is also cut off at the deadline here. On stdio the poll runs in a worker + # thread that a cancelled wait cannot stop, so it is given the same budget (budget_s) and + # ends by itself, releasing the connection. + remaining = max(0.0, deadline - time.time()) + if remaining <= 0: + continue + poll_params.pop("timeoutSeconds", None) + poll_params["timeout_seconds"] = max(1.0, remaining) try: - response = await send_with_unity_instance( - async_send_command_with_retry, - unity_instance, - tool_name, - poll_params, - user_id=user_id, + response = await asyncio.wait_for( + send_with_unity_instance( + async_send_command_with_retry, + unity_instance, + tool_name, + poll_params, + user_id=user_id, + budget_s=remaining, + ), + timeout=remaining, ) + except asyncio.TimeoutError: + # The deadline passed mid-poll: keep the last pending response and let the + # check above report the timeout. + continue except Exception as exc: # pragma: no cover - network/domain reload variability logger.debug(f"Polling {tool_name} failed, will retry: {exc}") # Back off modestly but stay responsive. @@ -275,12 +308,18 @@ def _normalize_response(self, response) -> MCPResponse: if isinstance(response, MCPResponse): return response if isinstance(response, dict): + data = response.get( + "data", response) if "data" not in response else response["data"] + hint = response.get("hint") + if isinstance(data, dict) and data.get("reason") == "unity_timeout": + # The tool may have run: its effects are to be checked, not the tool resubmitted. + hint = None return MCPResponse( success=response.get("success", True), message=response.get("message"), error=response.get("error"), - data=response.get( - "data", response) if "data" not in response else response["data"], + data=data, + hint=hint, ) success = True diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index 986781bed0..de42d716e7 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -27,6 +27,10 @@ FRAMED_MAX = 64 * 1024 * 1024 +class _DeadlineBeforeSend(TimeoutError): + """The command's deadline passed before its payload was written: nothing was sent, the socket is intact.""" + + @dataclass class UnityConnection: """Manages the socket connection to the Unity Editor.""" @@ -50,8 +54,11 @@ def _prepare_socket(self, sock: socket.socket) -> None: except OSError as exc: logger.debug(f"Unable to set TCP_NODELAY: {exc}") - def connect(self, connect_timeout: float | None = None) -> bool: - """Establish a connection to the Unity Editor.""" + def connect(self, connect_timeout: float | None = None, deadline: float | None = None) -> bool: + """Establish a connection to the Unity Editor. + + ``deadline`` (a monotonic() ceiling) caps both the connect and the handshake timeouts. + """ if self.sock: return True with self._conn_lock: @@ -62,6 +69,7 @@ def connect(self, connect_timeout: float | None = None) -> bool: if connect_timeout is None: connect_timeout = float( getattr(config, "connection_timeout", 1.0)) + connect_timeout = self._cap_to_deadline(connect_timeout, deadline) # We trust config.unity_host (default 127.0.0.1) but future improvements # could dynamically prefer 'localhost' depending on OS resolver behavior. self.sock = socket.create_connection( @@ -73,14 +81,17 @@ def connect(self, connect_timeout: float | None = None) -> bool: # Strict handshake: require FRAMING=1 try: require_framing = getattr(config, "require_framing", True) - handshake_timeout = float( - getattr(config, "handshake_timeout", 1.0)) + handshake_timeout = self._cap_to_deadline(float( + getattr(config, "handshake_timeout", 1.0)), deadline) self.sock.settimeout(handshake_timeout) buf = bytearray() - deadline = time.monotonic() + handshake_timeout - while time.monotonic() < deadline and len(buf) < 512: + # The earlier of the configured handshake timeout and the caller's deadline. + handshake_deadline = time.monotonic() + handshake_timeout + if deadline is not None: + handshake_deadline = min(deadline, handshake_deadline) + while time.monotonic() < handshake_deadline and len(buf) < 512: try: - chunk = self.sock.recv(256) + chunk = self._recv_by(self.sock, 256, handshake_deadline, handshake_timeout) if not chunk: break buf.extend(chunk) @@ -140,9 +151,13 @@ def _ensure_live_connection(self) -> None: """ if not self.sock: return - orig_blocking = None + orig_timeout = None + peeked = False try: - orig_blocking = self.sock.getblocking() + # setblocking(True) would clear the timeout, leaving the next send and receive unbounded: + # restore the timeout itself instead. + orig_timeout = self.sock.gettimeout() + peeked = True self.sock.setblocking(False) data = self.sock.recv(1, socket.MSG_PEEK) if not data: @@ -157,21 +172,35 @@ def _ensure_live_connection(self) -> None: pass self.sock = None finally: - if self.sock and orig_blocking is not None: - self.sock.setblocking(orig_blocking) - - def _read_exact(self, sock: socket.socket, count: int) -> bytes: + if self.sock and peeked: + self.sock.settimeout(orig_timeout) + + def _recv_by(self, sock: socket.socket, size: int, deadline: float | None, timeout: float | None) -> bytes: + """One recv that waits no longer than ``timeout`` nor past ``deadline`` (a monotonic() ceiling).""" + if deadline is not None: + if time.monotonic() >= deadline: + raise socket.timeout("Deadline reached while receiving") + sock.settimeout(self._cap_to_deadline(timeout or config.connection_timeout, deadline)) + return sock.recv(size) + + def _read_exact(self, sock: socket.socket, count: int, deadline: float | None = None, + timeout: float | None = None) -> bytes: data = bytearray() while len(data) < count: - chunk = sock.recv(count - len(data)) + chunk = self._recv_by(sock, count - len(data), deadline, timeout) if not chunk: raise ConnectionError( "Connection closed before reading expected bytes") data.extend(chunk) return bytes(data) - def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: - """Receive a complete response from Unity, handling chunked data.""" + def receive_full_response(self, sock, buffer_size=config.buffer_size, deadline: float | None = None) -> bytes: + """Receive a complete response from Unity, handling chunked data. + + With a ``deadline`` (a monotonic() ceiling), every read is capped to what is left of it, so a + response arriving in pieces cannot hold the connection past the deadline. + """ + timeout = sock.gettimeout() if self.use_framing: # Heartbeat semantics: the Unity editor emits zero-length frames while # a long-running command is still executing. We tolerate a bounded @@ -183,7 +212,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: heartbeat_count = 0 try: while True: - header = self._read_exact(sock, 8) + header = self._read_exact(sock, 8, deadline, timeout) payload_len = struct.unpack('>Q', header)[0] if payload_len == 0: heartbeat_count += 1 @@ -197,7 +226,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: if payload_len > FRAMED_MAX: raise ValueError( f"Invalid framed length: {payload_len}") - payload = self._read_exact(sock, payload_len) + payload = self._read_exact(sock, payload_len, deadline, timeout) logger.debug( f"Received framed response ({len(payload)} bytes)") return payload @@ -214,7 +243,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: # Respect the socket's currently configured timeout try: while True: - chunk = sock.recv(buffer_size) + chunk = self._recv_by(sock, buffer_size, deadline, timeout) if not chunk: if not chunks: raise Exception( @@ -267,6 +296,10 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: logger.error(f"Error during receive: {str(e)}") raise + def _lock_wait(self, deadline: float | None) -> float: + """How long to wait for the I/O lock: until the deadline (-1, without limit, when there is none).""" + return -1 if deadline is None else max(0.0, deadline - time.monotonic()) + def _cap_to_deadline(self, timeout: float, deadline: float | None, floor: float = 0.05) -> float: """Shrink a blocking timeout to whatever budget remains before the deadline.""" if deadline is None: @@ -341,8 +374,13 @@ def read_status_file(target_hash: str | None = None) -> dict | None: if status and (status.get('reloading') or status.get('reason') == 'reloading'): # Reload invalidates the socket; drop it under the I/O lock so this # close is serialized against the send/recv block, then reconnect next call. - with self._io_lock: - self.disconnect() + # Wait for the lock only while the budget lasts: if another command still + # holds it, that command owns the socket and the reload ends its exchange. + if self._io_lock.acquire(timeout=self._lock_wait(deadline)): + try: + self.disconnect() + finally: + self._io_lock.release() return MCPResponse( success=False, error="Unity is reloading; please retry", @@ -359,148 +397,163 @@ def read_status_file(target_hash: str | None = None) -> dict | None: raise TimeoutError( f"Command '{command_type}' exceeded total deadline of " f"{total_timeout:.1f}s (connection wedged or Unity unresponsive)") + # Commands on this connection are serialized, from the liveness check to the close after a + # failure, so no command changes the socket (or its timeout) while another is using it. Wait for + # another command's turn only as long as this one's budget lasts. + if not self._io_lock.acquire(timeout=self._lock_wait(deadline)): + raise _DeadlineBeforeSend( + f"Command '{command_type}' ran out of time waiting for the connection to Unity") + failure: Exception | None = None try: - # Discard stale sockets left over from a previous domain reload - # so we reconnect instead of writing to a dead connection. - self._ensure_live_connection() - # Ensure connected (handshake occurs within connect()) - t_conn_start = time.time() - if not self.sock and not self.connect(self._cap_to_deadline(config.connection_timeout, deadline)): - raise ConnectionError("Could not connect to Unity") - logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type) - - # Build payload - if command_type == 'ping': - payload = b'ping' - else: - payload = json.dumps({ - 'type': command_type, - 'params': params, - }).encode('utf-8') - - # Send/receive are serialized to protect the shared socket - with self._io_lock: + try: + # Discard stale sockets left over from a previous domain reload + # so we reconnect instead of writing to a dead connection. + self._ensure_live_connection() + # Ensure connected (handshake occurs within connect()) + t_conn_start = time.time() + if not self.sock and not self.connect(self._cap_to_deadline(config.connection_timeout, deadline), deadline=deadline): + raise ConnectionError("Could not connect to Unity") + logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type) + + # Build payload + if command_type == 'ping': + payload = b'ping' + else: + payload = json.dumps({ + 'type': command_type, + 'params': params, + }).encode('utf-8') + + if deadline is not None and time.monotonic() >= deadline: + raise _DeadlineBeforeSend( + f"Command '{command_type}' ran out of time before it was sent") mode = 'framed' if self.use_framing else 'legacy' with contextlib.suppress(Exception): logger.debug( f"send {len(payload)} bytes; mode={mode}; head={payload[:32].decode('utf-8', 'ignore')}") - t_send_start = time.time() - if self.use_framing: - header = struct.pack('>Q', len(payload)) - self.sock.sendall(header) - self.sock.sendall(payload) - else: - self.sock.sendall(payload) - logger.info("[TIMING-STDIO] sendall took %.3fs command=%s", time.time() - t_send_start, command_type) - - # Cap the receive timeout to the remaining command budget (and use a - # short timeout during retry bursts) so a wedged socket can't block - # past the deadline. - restore_timeout = None - recv_timeout = 1.0 if attempt > 0 else self.sock.gettimeout() + # Cap sends and receives to the remaining command budget (and use a short timeout + # during retry bursts) so a wedged socket can't block past the deadline. + original_timeout = self.sock.gettimeout() + io_timeout = 1.0 if attempt > 0 else original_timeout if deadline is not None: - recv_timeout = self._cap_to_deadline( - recv_timeout or config.connection_timeout, deadline) - if recv_timeout is not None and recv_timeout != self.sock.gettimeout(): - restore_timeout = self.sock.gettimeout() - self.sock.settimeout(recv_timeout) + io_timeout = self._cap_to_deadline( + io_timeout or config.connection_timeout, deadline) try: + self.sock.settimeout(io_timeout) + t_send_start = time.time() + if self.use_framing: + header = struct.pack('>Q', len(payload)) + self.sock.sendall(header) + self.sock.sendall(payload) + else: + self.sock.sendall(payload) + logger.info("[TIMING-STDIO] sendall took %.3fs command=%s", time.time() - t_send_start, command_type) + + if deadline is not None: + self.sock.settimeout(self._cap_to_deadline(io_timeout, deadline)) t_recv_start = time.time() - response_data = self.receive_full_response(self.sock) + response_data = self.receive_full_response(self.sock, deadline=deadline) logger.info("[TIMING-STDIO] receive took %.3fs command=%s len=%d", time.time() - t_recv_start, command_type, len(response_data)) with contextlib.suppress(Exception): logger.debug( f"recv {len(response_data)} bytes; mode={mode}") finally: - if restore_timeout is not None: - self.sock.settimeout(restore_timeout) + if self.sock: + self.sock.settimeout(original_timeout) - # Parse - if command_type == 'ping': - resp = json.loads(response_data.decode('utf-8')) - if resp.get('status') == 'success' and resp.get('result', {}).get('message') == 'pong': - return {"message": "pong"} - raise Exception("Ping unsuccessful") - - resp = json.loads(response_data.decode('utf-8')) - if resp.get('status') == 'error': - err = resp.get('error') or resp.get( - 'message', 'Unknown Unity error') - raise Exception(err) - return resp.get('result', {}) - except Exception as e: - logger.warning( - f"Unity communication attempt {attempt+1} failed: {e}") - try: - if self.sock: - self.sock.close() - finally: - self.sock = None + # Parse + if command_type == 'ping': + resp = json.loads(response_data.decode('utf-8')) + if resp.get('status') == 'success' and resp.get('result', {}).get('message') == 'pong': + return {"message": "pong"} + raise Exception("Ping unsuccessful") - # Re-discover the port for this specific instance - try: - stdio_port_registry.get_instances(force_refresh=True) - new_port: int | None = None - if self.instance_id: - # Try to rediscover the specific instance via shared registry - refreshed_instance = stdio_port_registry.get_instance( - self.instance_id) - if refreshed_instance and isinstance(refreshed_instance.port, int): - new_port = refreshed_instance.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; falling back to port scan", - ) + resp = json.loads(response_data.decode('utf-8')) + if resp.get('status') == 'error': + err = resp.get('error') or resp.get( + 'message', 'Unknown Unity error') + raise Exception(err) + return resp.get('result', {}) + except _DeadlineBeforeSend: + # Nothing was sent, and the socket is left as it was. + raise + except Exception as e: + failure = e + logger.warning( + f"Unity communication attempt {attempt+1} failed: {e}") + try: + if self.sock: + self.sock.close() + finally: + self.sock = None + finally: + self._io_lock.release() - # Fallback to registry default if instance-specific discovery failed - if new_port is None: - new_port = stdio_port_registry.get_port( - self.instance_id) - logger.info( - f"Using Unity port from stdio_port_registry: {new_port}") - - if new_port != self.port: - logger.info( - f"Unity port changed {self.port} -> {new_port}") - self.port = new_port - except Exception as de: - logger.debug(f"Port discovery failed: {de}") - - if attempt < attempts: - # Heartbeat-aware, jittered backoff - status = read_status_file(target_hash) - # Base exponential backoff - backoff = base_backoff * (2 ** attempt) - # Decorrelated jitter multiplier - jitter = random.uniform(0.1, 0.3) - - # Fast‑retry for transient socket failures - fast_error = isinstance( - e, (ConnectionRefusedError, ConnectionResetError, TimeoutError)) - if not fast_error: - try: - err_no = getattr(e, 'errno', None) - fast_error = err_no in ( - errno.ECONNREFUSED, errno.ECONNRESET, errno.ETIMEDOUT) - except Exception: - pass - - # Cap backoff depending on state - if status and status.get('reloading'): - # Domain reload can take 10-20s; use longer waits - cap = 5.0 - elif fast_error: - cap = 0.25 + e = failure + # Re-discover the port for this specific instance + try: + stdio_port_registry.get_instances(force_refresh=True) + new_port: int | None = None + if self.instance_id: + # Try to rediscover the specific instance via shared registry + refreshed_instance = stdio_port_registry.get_instance( + self.instance_id) + if refreshed_instance and isinstance(refreshed_instance.port, int): + new_port = refreshed_instance.port + logger.debug( + f"Rediscovered instance {self.instance_id} on port {new_port}") else: - cap = 3.0 + logger.warning( + f"Instance {self.instance_id} not found during reconnection; falling back to port scan", + ) + + # Fallback to registry default if instance-specific discovery failed + if new_port is None: + new_port = stdio_port_registry.get_port( + self.instance_id) + logger.info( + f"Using Unity port from stdio_port_registry: {new_port}") - sleep_s = min(cap, jitter * (2 ** attempt)) - sleep_s = self._cap_to_deadline(sleep_s, deadline, floor=0.0) - time.sleep(sleep_s) - continue - raise + if new_port != self.port: + logger.info( + f"Unity port changed {self.port} -> {new_port}") + self.port = new_port + except Exception as de: + logger.debug(f"Port discovery failed: {de}") + + if attempt < attempts: + # Heartbeat-aware, jittered backoff + status = read_status_file(target_hash) + # Base exponential backoff + backoff = base_backoff * (2 ** attempt) + # Decorrelated jitter multiplier + jitter = random.uniform(0.1, 0.3) + + # Fast‑retry for transient socket failures + fast_error = isinstance( + e, (ConnectionRefusedError, ConnectionResetError, TimeoutError)) + if not fast_error: + try: + err_no = getattr(e, 'errno', None) + fast_error = err_no in ( + errno.ECONNREFUSED, errno.ECONNRESET, errno.ETIMEDOUT) + except Exception: + pass + + # Cap backoff depending on state + if status and status.get('reloading'): + # Domain reload can take 10-20s; use longer waits + cap = 5.0 + elif fast_error: + cap = 0.25 + else: + cap = 3.0 + + sleep_s = min(cap, jitter * (2 ** attempt)) + sleep_s = self._cap_to_deadline(sleep_s, deadline, floor=0.0) + time.sleep(sleep_s) + continue + raise e # ----------------------------- @@ -670,13 +723,14 @@ def _resolve_instance_id(self, instance_identifier: str | None, instances: list[ f"Check mcpforunity://instances resource for all instances." ) - def get_connection(self, instance_identifier: str | None = None) -> UnityConnection: + def get_connection(self, instance_identifier: str | None = None, deadline: float | None = None) -> UnityConnection: """ Get or create a connection to a Unity instance. Args: instance_identifier: Optional identifier (name, hash, name@hash, etc.) If None, uses default or most recent instance + deadline: Optional monotonic() ceiling for establishing a new connection Returns: UnityConnection to the specified instance @@ -696,7 +750,7 @@ def get_connection(self, instance_identifier: str | None = None) -> UnityConnect logger.info( f"Creating new connection to Unity instance: {target.id} (port {target.port})") conn = UnityConnection(port=target.port, instance_id=target.id) - if not conn.connect(): + if not conn.connect(deadline=deadline): raise ConnectionError( f"Failed to connect to Unity instance '{target.id}' on port {target.port}. " f"Ensure the Unity Editor is running." @@ -749,7 +803,7 @@ def get_unity_connection_pool() -> UnityConnectionPool: # Backwards compatibility: keep old single-connection function -def get_unity_connection(instance_identifier: str | None = None) -> UnityConnection: +def get_unity_connection(instance_identifier: str | None = None, deadline: float | None = None) -> UnityConnection: """Retrieve or establish a Unity connection. Args: @@ -762,7 +816,7 @@ def get_unity_connection(instance_identifier: str | None = None) -> UnityConnect Note: This function now uses the connection pool internally. """ pool = get_unity_connection_pool() - return pool.get_connection(instance_identifier) + return pool.get_connection(instance_identifier, deadline=deadline) # ----------------------------- @@ -818,7 +872,8 @@ def send_command_with_retry( instance_id: str | None = None, max_retries: int | None = None, retry_ms: int | None = None, - retry_on_reload: bool = True + retry_on_reload: bool = True, + budget_s: float | None = None, ) -> dict[str, Any] | MCPResponse: """Send a command to a Unity instance, waiting politely through Unity reloads. @@ -830,6 +885,8 @@ def send_command_with_retry( 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) + budget_s: Optional cap in seconds on the whole call, connection retries and reload + waits included; it can only shorten config.command_total_timeout Returns: Response dictionary or MCPResponse from Unity @@ -839,9 +896,23 @@ def send_command_with_retry( """ t_retry_start = time.time() logger.info("[TIMING-STDIO] send_command_with_retry START command=%s", command_type) + # The budget starts now, so finding and connecting to the instance spends it too. + total_timeout = max(0.0, float(getattr(config, "command_total_timeout", 90.0))) + deadline = time.monotonic() + total_timeout if total_timeout > 0 else None + if budget_s is not None: + capped = time.monotonic() + max(0.0, budget_s) + deadline = capped if deadline is None else min(deadline, capped) t_get_conn = time.time() - conn = get_unity_connection(instance_id) + conn = get_unity_connection(instance_id, deadline=deadline) logger.info("[TIMING-STDIO] get_unity_connection took %.3fs command=%s", time.time() - t_get_conn, command_type) + if deadline is not None and time.monotonic() >= deadline: + # The caller has given up by now (a status poll past its deadline): do not send a command it no + # longer waits for. + return MCPResponse( + success=False, + error=f"No time left to send '{command_type}' after connecting to Unity", + hint="retry", + ) if max_retries is None: max_retries = getattr(config, "reload_max_retries", 40) if retry_ms is None: @@ -868,9 +939,6 @@ def send_command_with_retry( # Clamp to [0, 20] to prevent misconfiguration from causing excessive waits max_wait_s = max(0.0, min(max_wait_s, 20.0)) - total_timeout = max(0.0, float(getattr(config, "command_total_timeout", 90.0))) - deadline = time.monotonic() + total_timeout if total_timeout > 0 else None - # 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 @@ -952,7 +1020,8 @@ async def async_send_command_with_retry( loop=None, max_retries: int | None = None, retry_ms: int | None = None, - retry_on_reload: bool = True + retry_on_reload: bool = True, + budget_s: float | None = None, ) -> dict[str, Any] | MCPResponse: """Async wrapper that runs the blocking retry helper in a thread pool. @@ -964,6 +1033,7 @@ async def async_send_command_with_retry( 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 + budget_s: Optional cap in seconds on the whole call (see send_command_with_retry) Returns: Response dictionary or MCPResponse on error @@ -976,7 +1046,7 @@ async def async_send_command_with_retry( None, lambda: send_command_with_retry( command_type, params, instance_id=instance_id, max_retries=max_retries, - retry_ms=retry_ms, retry_on_reload=retry_on_reload), + retry_ms=retry_ms, retry_on_reload=retry_on_reload, budget_s=budget_s), ) # After a successful command, check if the connection was freshly diff --git a/Server/src/transport/unity_transport.py b/Server/src/transport/unity_transport.py index 9d07f4714c..57149c3d9c 100644 --- a/Server/src/transport/unity_transport.py +++ b/Server/src/transport/unity_transport.py @@ -1,6 +1,7 @@ """Transport helpers for routing commands to Unity.""" from __future__ import annotations +import asyncio import logging from typing import Awaitable, Callable, TypeVar @@ -101,11 +102,25 @@ async def send_with_unity_instance( except Exception as exc: # NOTE: asyncio.TimeoutError has an empty str() by default, which is confusing for clients. err = str(exc) or f"{type(exc).__name__}" + data: dict | None = None + if isinstance(exc, (asyncio.TimeoutError, TimeoutError)): + # A bare "TimeoutError" reads as noise and agents plough on as if the + # command ran. Spell out what a timeout means and what to do about it. + err = ( + f"Unity did not respond to '{command_type}' within its time budget. " + "The editor is likely compiling, reloading, importing assets, or still " + "running the operation. The command may or may not have executed — do NOT " + "assume it ran; verify its effects (or re-run it) once " + "mcpforunity://editor/state reports advice.ready_for_tools=true. " + "For long-running tools, pass timeout_seconds in the tool parameters " + "(honored up to 3600)." + ) + data = {"reason": "unity_timeout"} # Fail fast with a retry hint instead of hanging for COMMAND_TIMEOUT. # The client can decide whether retrying is appropriate for the command. return normalize_unity_response( MCPResponse(success=False, error=err, - hint="retry").model_dump() + hint="retry", data=data).model_dump() ) if unity_instance: diff --git a/Server/tests/integration/test_connection_deadline.py b/Server/tests/integration/test_connection_deadline.py index dd6b280e43..cfc1ff76f1 100644 --- a/Server/tests/integration/test_connection_deadline.py +++ b/Server/tests/integration/test_connection_deadline.py @@ -91,7 +91,7 @@ def test_send_command_against_wedged_socket_is_bounded(silent_bridge, monkeypatc def test_send_command_with_retry_is_bounded_by_deadline(silent_bridge, tmp_path, monkeypatch): """The reload-wait loop honours command_total_timeout, not just max_wait_s.""" conn = UnityConnection(port=silent_bridge, instance_id="Repro@deadbeef") - monkeypatch.setattr(uc, "get_unity_connection", lambda instance_id=None: conn) + monkeypatch.setattr(uc, "get_unity_connection", lambda instance_id=None, deadline=None: conn) _write_reloading_status(tmp_path) start = time.monotonic() @@ -143,7 +143,7 @@ def discover_instances() -> list[UnityInstanceInfo]: attempted_ports: list[int] = [] ports_before_backoff: list[int] = [] - def fail_connect(connect_timeout: float | None = None) -> bool: + def fail_connect(connect_timeout: float | None = None, deadline: float | None = None) -> bool: attempted_ports.append(conn.port) return False diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py new file mode 100644 index 0000000000..1252e16040 --- /dev/null +++ b/Server/tests/test_custom_tool_timeouts.py @@ -0,0 +1,479 @@ +"""Timeouts on custom tools must be long-budget, loud, and impossible to misread. + +Real-world failure (2026-08-11): a long-running project tool (image comparison) +exceeded the hub's default 30s command budget and the caller got the five words +{"success":false,"error":"TimeoutError","hint":null} — after which the agent +carried on reading stale on-disk artifacts as if the tool had run. Three fixes: +custom tools default to a 300s budget, _normalize_response no longer drops the +hint, and transport-level timeouts explain that the command may not have run. +""" + +import asyncio + +import pytest + +from core.config import config +from models.models import MCPResponse, ToolDefinitionModel +from services.custom_tool_service import CustomToolService + + +class _NoopMcp: + def custom_route(self, _path, methods=None): # noqa: ARG002 + def _decorator(fn): + return fn + + return _decorator + + def tool(self, name=None, description=None): # noqa: ARG002 + def _decorator(fn): + return fn + + return _decorator + + +def _make_service() -> CustomToolService: + return CustomToolService(_NoopMcp(), project_scoped_tools=True) + + +def test_normalize_response_preserves_hint(): + service = _make_service() + result = service._normalize_response( + {"success": False, "error": "boom", "hint": "retry", "data": {"reason": "x"}} + ) + assert isinstance(result, MCPResponse) + assert result.hint == "retry" + assert result.data == {"reason": "x"} + + +@pytest.mark.asyncio +async def test_execute_tool_defaults_long_timeout(monkeypatch): + service = _make_service() + definition = ToolDefinitionModel(name="long_running_tool", description="long tool") + service._register_tool("proj-1", definition) + + seen_params: dict = {} + + async def fake_send(_send_fn, _instance, _command, params, **kwargs): # noqa: ARG001 + seen_params.update(params) + return {"success": True, "message": "ok"} + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", fake_send) + + result = await service.execute_tool("proj-1", "long_running_tool", "Playground@abc", {}) + assert result.success is True + assert seen_params["timeout_seconds"] == 300 + + # An explicit caller value must win over the default. + seen_params.clear() + await service.execute_tool( + "proj-1", "long_running_tool", "Playground@abc", {"timeout_seconds": 900}) + assert seen_params["timeout_seconds"] == 900 + + # So must the camelCase spelling the hub also honors: no snake_case default beside it. + seen_params.clear() + await service.execute_tool( + "proj-1", "long_running_tool", "Playground@abc", {"timeoutSeconds": 1200}) + assert seen_params == {"timeoutSeconds": 1200} + + +@pytest.mark.asyncio +async def test_polling_deadline_does_not_suggest_resubmitting(monkeypatch): + """At the polling deadline the operation may still run: say so, and do not hint a retry.""" + service = _make_service() + + async def still_running(_send_fn, _instance, _command, params, **kwargs): # noqa: ARG001 + return {"_mcp_status": "pending", "message": "working"} + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", still_running) + + result = await service._poll_until_complete( + "long_running_tool", "Playground@abc", {}, {"_mcp_status": "pending"}, "status", + max_poll_seconds=0.01) + + assert result.success is False + assert "may still be running" in result.message + assert result.hint is None + + +@pytest.mark.asyncio +async def test_transport_timeout_error_is_self_explanatory(monkeypatch): + from transport import unity_transport + + monkeypatch.setattr(config, "transport_mode", "http") + + async def raise_timeout(*args, **kwargs): # noqa: ARG001 + raise asyncio.TimeoutError() + + monkeypatch.setattr( + unity_transport.PluginHub, "send_command_for_instance", raise_timeout) + + async def unused_send_fn(*args, **kwargs): # noqa: ARG001 + raise AssertionError("stdio path should not be used in http mode") + + response = await unity_transport.send_with_unity_instance( + unused_send_fn, "Playground@abc", "long_running_tool", {}) + + assert response["success"] is False + assert "may or may not have executed" in response["error"] + assert "timeout_seconds" in response["error"] + assert response["hint"] == "retry" + assert response["data"] == {"reason": "unity_timeout"} + + +@pytest.mark.asyncio +async def test_each_poll_waits_no_longer_than_the_remaining_budget(monkeypatch): + """A poll must not inherit the tool's 300 s command budget past the polling deadline.""" + service = _make_service() + seen: list[dict] = [] + + async def complete(_send_fn, _instance, _command, params, **kwargs): # noqa: ARG001 + seen.append(dict(params)) + return {"_mcp_status": "complete", "success": True, "message": "done"} + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", complete) + + result = await service._poll_until_complete( + "long_running_tool", "Playground@abc", {"timeout_seconds": 300, "timeoutSeconds": 300}, + {"_mcp_status": "pending", "_mcp_poll_interval": 0.01}, "status", max_poll_seconds=5) + + assert result.success is True + assert seen and 0 < seen[0]["timeout_seconds"] <= 5 + assert "timeoutSeconds" not in seen[0] + assert seen[0]["action"] == "status" + + +@pytest.mark.asyncio +async def test_a_hung_poll_is_cut_off_at_the_polling_deadline(monkeypatch): + """The hub waits at least 30 s for any command; a poll must still end at the polling deadline.""" + service = _make_service() + + async def hang(_send_fn, _instance, _command, _params, **kwargs): # noqa: ARG001 + await asyncio.sleep(30) + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", hang) + + loop = asyncio.get_running_loop() + started = loop.time() + result = await service._poll_until_complete( + "long_running_tool", "Playground@abc", {}, + {"_mcp_status": "pending", "_mcp_poll_interval": 0.01}, "status", max_poll_seconds=1) + elapsed = loop.time() - started + + assert result.success is False + assert "Timeout waiting for long_running_tool" in result.message + assert elapsed < 2.0 + + +@pytest.mark.asyncio +async def test_a_stdio_poll_stops_its_worker_thread_at_the_polling_deadline(monkeypatch): + """On stdio the poll runs in an executor thread that a cancelled wait cannot stop. It must be given the + remaining budget, so the thread itself ends at the polling deadline and releases the connection.""" + import time + + import transport.legacy.unity_connection as unity_connection + + calls: list[dict] = [] + + class _BlockingConnection: + def send_command(self, command_type, params, max_attempts=None, deadline=None): # noqa: ARG002 + call = {"deadline": deadline, "ended": None} + calls.append(call) + started = time.monotonic() + # Block like a wedged socket would, until the deadline the caller allows (5 s at most here). + time.sleep(min(5.0, max(0.0, (deadline or started + 5.0) - started))) + call["ended"] = time.monotonic() + return {"_mcp_status": "pending", "_mcp_poll_interval": 0.01} + + monkeypatch.setattr(config, "transport_mode", "stdio") + monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None, deadline=None: _BlockingConnection()) + + service = _make_service() + started = time.monotonic() + result = await service._poll_until_complete( + "long_running_tool", None, {}, + {"_mcp_status": "pending", "_mcp_poll_interval": 0.01}, "status", max_poll_seconds=1) + await asyncio.sleep(0.2) # let a thread still blocked past the deadline (the bug) record itself + + assert result.success is False + assert calls, "the poll must have reached the stdio connection" + assert all(c["deadline"] is not None and c["deadline"] <= started + 1.5 for c in calls) + assert all(c["ended"] is not None and c["ended"] <= started + 1.5 for c in calls), "a poll thread outlived the deadline" + + +def test_budget_only_shortens_the_configured_command_deadline(monkeypatch): + import time + + import transport.legacy.unity_connection as unity_connection + + deadlines: list[float] = [] + + class _Connection: + def send_command(self, command_type, params, max_attempts=None, deadline=None): # noqa: ARG002 + deadlines.append(deadline - time.monotonic()) + return {"success": True} + + monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None, deadline=None: _Connection()) + monkeypatch.setattr(config, "command_total_timeout", 90.0) + + unity_connection.send_command_with_retry("x", {}, budget_s=2.0) + unity_connection.send_command_with_retry("x", {}, budget_s=500.0) + unity_connection.send_command_with_retry("x", {}) + + assert deadlines[0] <= 2.0 + assert 85.0 < deadlines[1] <= 90.0 + assert 85.0 < deadlines[2] <= 90.0 + + +def test_finding_the_instance_spends_the_budget(monkeypatch): + """The budget starts before the connection lookup: a slow lookup leaves less for the command, and a + lookup that spends all of it means the command is not sent at all.""" + import time + + import transport.legacy.unity_connection as unity_connection + + sent: list[float] = [] + lookup_deadlines: list[float | None] = [] + + class _Connection: + def send_command(self, command_type, params, max_attempts=None, deadline=None): # noqa: ARG002 + sent.append(deadline - time.monotonic()) + return {"success": True} + + def slow_lookup(instance_id=None, deadline=None): # noqa: ARG001 + lookup_deadlines.append(deadline) + time.sleep(0.3) + return _Connection() + + monkeypatch.setattr(unity_connection, "get_unity_connection", slow_lookup) + monkeypatch.setattr(config, "command_total_timeout", 90.0) + + started = time.monotonic() + unity_connection.send_command_with_retry("x", {}, budget_s=1.0) + assert lookup_deadlines[0] is not None and lookup_deadlines[0] <= started + 1.05 + assert sent and sent[0] <= 0.75 + + result = unity_connection.send_command_with_retry("x", {}, budget_s=0.2) + assert len(sent) == 1, "a command whose budget the lookup spent must not be sent" + assert result.success is False + + +def test_connect_caps_its_timeouts_to_the_deadline(monkeypatch): + import socket + import time + + import transport.legacy.unity_connection as unity_connection + + timeouts: list[float] = [] + + def refuse(address, timeout=None): # noqa: ARG001 + timeouts.append(timeout) + raise ConnectionRefusedError("no editor here") + + monkeypatch.setattr(socket, "create_connection", refuse) + monkeypatch.setattr(config, "connection_timeout", 30.0) + + conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") + assert conn.connect(deadline=time.monotonic() + 0.5) is False + assert timeouts and timeouts[0] <= 0.5 + + +def test_a_timed_out_custom_tool_carries_no_retry_hint(): + """A timeout may follow a command that ran: the caller is told to check, not to resubmit.""" + service = _make_service() + result = service._normalize_response( + {"success": False, "error": "late", "hint": "retry", "data": {"reason": "unity_timeout"}}) + assert result.hint is None + assert result.data == {"reason": "unity_timeout"} + + +def _call_in_thread(fn, limit_s): + """Run fn in a thread; return (elapsed, outcome) or fail if it is still blocked after limit_s.""" + import threading + import time + + outcome: dict = {} + + def run(): + try: + outcome["result"] = fn() + except Exception as exc: # noqa: BLE001 + outcome["error"] = exc + + started = time.monotonic() + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(limit_s) + assert not worker.is_alive(), f"still blocked after {limit_s} s" + return time.monotonic() - started, outcome + + +def _stdio_connection(monkeypatch, tmp_path): + import socket + + import transport.legacy.unity_connection as unity_connection + + # No status files: the reload preflight must not see a real editor's. + monkeypatch.setattr(unity_connection.Path, "home", lambda: tmp_path) + monkeypatch.setattr(unity_connection.stdio_port_registry, "get_instances", lambda **_: []) + monkeypatch.setattr(unity_connection.stdio_port_registry, "get_instance", lambda *_: None) + monkeypatch.setattr(unity_connection.stdio_port_registry, "get_port", lambda *_: 1) + ours, unity = socket.socketpair() + ours.settimeout(30.0) + conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") + conn.sock = ours + conn.use_framing = True + return conn, unity + + +def test_a_command_waits_for_the_connection_only_while_its_budget_lasts(monkeypatch, tmp_path): + import time + + conn, unity = _stdio_connection(monkeypatch, tmp_path) + unity.setblocking(False) + conn._io_lock.acquire() # another command is using the connection + try: + elapsed, outcome = _call_in_thread( + lambda: conn.send_command("x", {}, max_attempts=0, deadline=time.monotonic() + 0.3), 3.0) + finally: + conn._io_lock.release() + assert isinstance(outcome.get("error"), TimeoutError) + assert elapsed < 1.0 + assert conn.sock is not None, "the other command's socket must be left alone" + try: + sent = unity.recv(64) + except BlockingIOError: + sent = b"" + assert sent == b"", "a command whose budget ran out waiting must not be sent" + + +def test_a_response_in_pieces_is_cut_off_at_the_deadline(monkeypatch, tmp_path): + import struct + import threading + import time + + conn, unity = _stdio_connection(monkeypatch, tmp_path) + stop = threading.Event() + + def trickle(): + unity.sendall(struct.pack(">Q", 1000)) + while not stop.is_set(): # one byte every 0.2 s, never the whole payload + try: + unity.sendall(b"x") + except OSError: + return + time.sleep(0.2) + + threading.Thread(target=trickle, daemon=True).start() + try: + elapsed, outcome = _call_in_thread( + lambda: conn.receive_full_response(conn.sock, deadline=time.monotonic() + 0.5), 5.0) + finally: + stop.set() + assert isinstance(outcome.get("error"), TimeoutError) + assert elapsed < 1.0 + + +def test_a_send_unity_does_not_read_is_cut_off_at_the_deadline(monkeypatch, tmp_path): + import time + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + big = {"blob": "x" * (8 * 1024 * 1024)} # far beyond the socket buffers; Unity reads none of it + elapsed, outcome = _call_in_thread( + lambda: conn.send_command("x", big, max_attempts=0, deadline=time.monotonic() + 0.5), 5.0) + assert isinstance(outcome.get("error"), TimeoutError) + assert elapsed < 1.5 + + +def test_the_liveness_check_keeps_the_socket_timeout(monkeypatch, tmp_path): + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + conn._ensure_live_connection() + assert conn.sock is not None + assert conn.sock.gettimeout() == 30.0 + + +def test_a_reconnect_handshake_spends_the_command_budget(monkeypatch, tmp_path): + import time + + import transport.legacy.unity_connection as unity_connection + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + conn.sock = None # the stale socket was dropped: send_command reconnects + seen: list = [] + + def fake_connect(self, connect_timeout=None, deadline=None): # noqa: ARG001 + seen.append(deadline) + return False + + monkeypatch.setattr(unity_connection.UnityConnection, "connect", fake_connect) + deadline = time.monotonic() + 0.5 + _call_in_thread(lambda: conn.send_command("x", {}, max_attempts=0, deadline=deadline), 3.0) + assert seen == [deadline] + + +def test_the_reload_preflight_waits_for_the_connection_only_while_the_budget_lasts(monkeypatch, tmp_path): + import json + import time + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + status_dir = tmp_path / ".unity-mcp" + status_dir.mkdir() + (status_dir / "unity-mcp-status-abc.json").write_text(json.dumps({"reloading": True})) + conn._io_lock.acquire() # another command is using the connection + try: + elapsed, outcome = _call_in_thread( + lambda: conn.send_command("x", {}, max_attempts=0, deadline=time.monotonic() + 0.3), 3.0) + finally: + conn._io_lock.release() + assert elapsed < 1.0 + assert outcome["result"].hint == "retry" + assert conn.sock is not None, "the socket of the command holding the lock is left alone" + + +def test_a_waiting_command_leaves_the_socket_alone_until_its_turn(monkeypatch, tmp_path): + """The liveness check reads and resets the socket's timeout: done while another command holds the + connection, it could put back a longer timeout under that command's send.""" + import time + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + checked: list[bool] = [] + monkeypatch.setattr(conn, "_ensure_live_connection", lambda: checked.append(True)) + conn._io_lock.acquire() # another command is using the connection + try: + _call_in_thread( + lambda: conn.send_command("x", {}, max_attempts=0, deadline=time.monotonic() + 0.3), 3.0) + finally: + conn._io_lock.release() + assert checked == [], "the liveness check ran while another command held the connection" + + +def test_the_handshake_ends_at_the_earlier_of_its_timeout_and_the_deadline(monkeypatch): + import socket + import threading + import time + + import transport.legacy.unity_connection as unity_connection + + ours, unity = socket.socketpair() + stop = threading.Event() + + def trickle(): # a byte every 0.45 s, never the end of the handshake line + while not stop.is_set(): + try: + unity.sendall(b"F") + except OSError: + return + time.sleep(0.45) + + monkeypatch.setattr(socket, "create_connection", lambda *_a, **_k: ours) + monkeypatch.setattr(config, "handshake_timeout", 5.0) + threading.Thread(target=trickle, daemon=True).start() + conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") + try: + elapsed, _ = _call_in_thread(lambda: conn.connect(deadline=time.monotonic() + 0.5), 5.0) + finally: + stop.set() + assert elapsed < 0.75